diff --git a/csp/adapters/perspective.py b/csp/adapters/perspective.py index b5a6acd6c..4abfbb23c 100644 --- a/csp/adapters/perspective.py +++ b/csp/adapters/perspective.py @@ -1,5 +1,6 @@ import threading from datetime import timedelta +from enum import IntEnum from typing import Dict, Optional, Union from perspective import Table as Table_, View as View_ @@ -235,7 +236,7 @@ def _instantiate(self): for table_name, table in self._tables.items(): schema = { - k: v.tstype.typ if not issubclass(v.tstype.typ, csp.Enum) else str for k, v in table.columns.items() + k: v.tstype.typ if not issubclass(v.tstype.typ, IntEnum) else str for k, v in table.columns.items() } if _PERSPECTIVE_3: psp_type_map = perspective_type_map() diff --git a/csp/impl/struct.py b/csp/impl/struct.py index 50540d365..b9e734e59 100644 --- a/csp/impl/struct.py +++ b/csp/impl/struct.py @@ -1,6 +1,7 @@ import io import typing from copy import deepcopy +from enum import IntEnum import ruamel.yaml from deprecated import deprecated @@ -224,7 +225,7 @@ def _obj_to_python(cls, obj): isinstance(obj, (list, tuple, set)) or type(obj).__name__ == "FastList" ): # hack for FastList that is not a list return type(obj)(cls._obj_to_python(v) for v in obj) - elif isinstance(obj, csp.Enum): + elif isinstance(obj, IntEnum): # any IntEnum is a native csp enum, not only csp.Enum return obj.name # handled in _obj_from_python else: return obj diff --git a/csp/tests/adapters/test_perspective.py b/csp/tests/adapters/test_perspective.py index 9bdf8fd00..bc8dbfb5b 100644 --- a/csp/tests/adapters/test_perspective.py +++ b/csp/tests/adapters/test_perspective.py @@ -1,5 +1,6 @@ import unittest from datetime import date, datetime, timedelta +from enum import IntEnum import csp from csp.utils.datetime import utc_now @@ -14,19 +15,30 @@ HAS_PERSPECTIVE = False +class MyIntEnum(IntEnum): + A = 1 + B = 2 + + class MyStruct(csp.Struct): my_str: str my_float: float my_bool: bool my_date: date my_datetime: datetime + my_enum: MyIntEnum def my_graph(output={}): adapter = PerspectiveAdapter(8000) table = adapter.create_table("Test") data = MyStruct( - my_str="foo", my_float=1.0, my_bool=False, my_date=date(2020, 1, 1), my_datetime=datetime(2020, 1, 1) + my_str="foo", + my_float=1.0, + my_bool=False, + my_date=date(2020, 1, 1), + my_datetime=datetime(2020, 1, 1), + my_enum=MyIntEnum.B, ) table.publish(csp.unroll(csp.const([data, data]))) output["table"] = table diff --git a/csp/tests/impl/test_struct.py b/csp/tests/impl/test_struct.py index 8bd0dc1ec..ddededf27 100644 --- a/csp/tests/impl/test_struct.py +++ b/csp/tests/impl/test_struct.py @@ -1555,6 +1555,24 @@ def postprocess_to_dict(self, obj): } self.assertEqual(test_struct.to_dict(), result_dict) + def test_to_dict_depr_int_enum(self): + class MyCspEnum(csp.Enum): + A = 1 + B = 2 + + class MyIntEnum(enum.IntEnum): + C = 3 + D = 4 + + class MyStruct(csp.Struct): + a: MyCspEnum + b: MyIntEnum + c: List[MyIntEnum] + + test_struct = MyStruct(a=MyCspEnum.B, b=MyIntEnum.C, c=[MyIntEnum.D]) + self.assertEqual(test_struct.to_dict_depr(), {"a": "B", "b": "C", "c": ["D"]}) + self.assertEqual(test_struct.to_dict_depr(), test_struct.to_dict()) + def test_to_dict_preserve_enums(self): class MyEnum(csp.Enum): A = 1