Merge pull request #322 from Point72/investigate/csp-pr743-intenum-compat-v3

Support CSP IntEnum compatibility

Author: timkpaine

csp_gateway/server/gateway/csp/channels.py

@@ -2,7 +2,7 @@
 from collections import defaultdict, deque
 from contextlib import contextmanager
 from datetime import datetime
-from enum import Enum
+from enum import Enum as PyEnum
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -14,7 +14,7 @@
 )
 
 import csp
-from csp import ts
+from csp import Enum as CspEnum, ts
 from csp.impl.genericpushadapter import GenericPushAdapter
 from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from csp.impl.types.tstype import TsType, isTsType
@@ -75,13 +75,18 @@ def _normalize_keyby(keyby: str | tuple[str, ...] | list) -> tuple[str, ...]:
     return (keyby,)
 
 
+def _has_indexer(indexer) -> bool:
+    # Preserve the legacy empty-string sentinel while accepting zero-valued keys.
+    return indexer not in (None, "")
+
+
 class _SnapshotModelBaseClass(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", coerce_numbers_to_str=True)
 
 
 def _recursive_remove_enums(vals_dict):
     for k, v in list(vals_dict.items()):
-        is_enum_key = isinstance(k, Enum)
+        is_enum_key = isinstance(k, (PyEnum, CspEnum))
         is_dict_value = isinstance(v, dict)
         if is_enum_key:
             v = vals_dict.pop(k)
@@ -375,9 +380,9 @@ def _add_field_to_graph(
             self._modules_connections_graph[field] = {"getters": [], "setters": []}
 
         if isinstance(module, str):
-            name = f"{module}{f'<{indexer}>' if indexer else ''}"
+            name = f"{module}{f'<{indexer}>' if _has_indexer(indexer) else ''}"
         else:
-            name = f"{module.__class__.__name__}{f'<{indexer}>' if indexer else ''}"
+            name = f"{module.__class__.__name__}{f'<{indexer}>' if _has_indexer(indexer) else ''}"
 
         if setting:
             if name not in self._modules_connections_graph[field]["setters"]:
@@ -400,7 +405,7 @@ def _keys_for_channel(self, field: str) -> dict[Any, _NONE_TYPE] | None:
         if is_dict_basket(tstype):
             # get type of key in basket
             basket_key_type = get_dict_basket_key_type(tstype)
-            if issubclass(basket_key_type, Enum):
+            if issubclass(basket_key_type, (PyEnum, CspEnum)):
                 return {e: None for e in basket_key_type}
             else:
                 return self._dynamic_keys.get(field, {})
@@ -727,7 +732,7 @@ def get_channel(
             # now set field to the delayed edge
             setattr(self, field, self._delayed_channels[field])
 
-        if is_dict_basket(tstype) and indexer:
+        if is_dict_basket(tstype) and _has_indexer(indexer):
             # if using an indexer, return that edge (raise if not recognized)
             if indexer not in self._keys_for_channel(field):
                 raise GatewayException(
@@ -770,11 +775,11 @@ def set_channel(
             gateway_tstype = tstype
 
         # validate arguments
-        if _is_dict_basket and isinstance(edge, Edge) and not indexer:
+        if _is_dict_basket and isinstance(edge, Edge) and not _has_indexer(indexer):
             # if its a dict basket and you set an edge, you need to provide an indexer
             raise GatewayException(f"Field `{field}` refers to a dict basket, and you have provided an edge {edge} but not an indexer")
 
-        if _is_dict_basket and isinstance(edge, dict) and indexer:
+        if _is_dict_basket and isinstance(edge, dict) and _has_indexer(indexer):
             # if its a dict basket and you set a dict, you should not provide an indexer
             raise GatewayException(f"Field `{field}` refers to a dict basket, and you have provided an edge basket but also an indexer {indexer}")
 
@@ -786,7 +791,7 @@ def set_channel(
             # must provide an edge type
             raise TypeError(f"Edge provided for field `{field}` is not an `Edge` instance {edge}")
 
-        if not _is_dict_basket and indexer:
+        if not _is_dict_basket and _has_indexer(indexer):
             # don't provide an indexer for non-dict basket
             raise GatewayException(f"Indexer provided for field `{field}` but it is not a basket instance")
 
@@ -967,7 +972,7 @@ def _set_last(self, field: str, indexer: str | int | None = None) -> None:
         tstype = self.get_outer_type(field)
 
         # First ensure edge is constructed
-        if indexer:
+        if _has_indexer(indexer):
             edge = self.get_channel(field, indexer)
         else:
             edge = self.get_channel(field)
@@ -990,7 +995,7 @@ def _set_last(self, field: str, indexer: str | int | None = None) -> None:
         trigger = ConcurrentFutureAdapter(name=f"RequestLast<{edge_type_name}>")
 
         if is_dict_basket(tstype):
-            if indexer:
+            if _has_indexer(indexer):
                 named_on_request_node_dict_basket(f"QueryLast<Basket<{edge_type_name}>>")({indexer: edge}, trigger.out())
             else:
                 named_on_request_node_dict_basket(f"QueryLast<Basket<{edge_type_name}>>")(edge, trigger.out())
@@ -1009,7 +1014,7 @@ def _set_next(self, field: str, indexer: str | int | None = None) -> None:
         tstype = self.get_outer_type(field)
 
         # First ensure edge is constructed
-        if indexer:
+        if _has_indexer(indexer):
             edge = self.get_channel(field, indexer)
         else:
             edge = self.get_channel(field)
@@ -1032,7 +1037,7 @@ def _set_next(self, field: str, indexer: str | int | None = None) -> None:
         trigger = ConcurrentFutureAdapter(name=f"RequestLast<{edge_type_name}>")
 
         if is_dict_basket(tstype):
-            if indexer:
+            if _has_indexer(indexer):
                 named_wait_for_next_node_dict_basket(f"QueryNext<Basket<{edge_type_name}>>")({indexer: edge}, trigger.out())
             else:
                 named_wait_for_next_node_dict_basket(f"QueryNext<Basket<{edge_type_name}>>")(edge, trigger.out())
@@ -1067,7 +1072,7 @@ def add_send_channel(self, field: str, indexer: str | int | None = None) -> None
         tstype = self.get_outer_type(field)
 
         if is_dict_basket(tstype):
-            if indexer:
+            if _has_indexer(indexer):
                 # wire in now
                 tstype = get_dict_basket_value_type(tstype)
                 self._send_channels[field, indexer] = GenericPushAdapter(tstype, name=f"manual_{field}[{indexer}]")
@@ -1081,7 +1086,7 @@ def add_send_channel(self, field: str, indexer: str | int | None = None) -> None
             tstype = tstype.typ
             self._send_channels[field, indexer] = GenericPushAdapter(tstype, name=f"manual_{field}")
 
-    def _add_send_channel_dict_basket(self, field: str, keys: list[str] | Enum) -> None:
+    def _add_send_channel_dict_basket(self, field: str, keys: list[str] | type[PyEnum] | type[CspEnum]) -> None:
         # NOTE: Do not call this directly, it is used in the factory finalization
 
         # get type of edge
@@ -1118,7 +1123,7 @@ def last(self, field: str, indexer: str | int | None = None, *, timeout=None) ->
 
         # wait for result
         result = future.result(timeout=timeout)
-        if indexer:
+        if _has_indexer(indexer):
             return result.get(indexer)
         return result
 
@@ -1134,7 +1139,7 @@ def next(self, field: str, indexer: str | int | None = None, *, timeout=None) ->
 
         # wait for result
         result = future.result(timeout=timeout)
-        if indexer:
+        if _has_indexer(indexer):
             return result.get(indexer)
         return result
 
@@ -1331,7 +1336,7 @@ def stage_lookup(
     def _check(self, field: str, where: dict, kind: str, indexer: str | int | None = None) -> None:
         if (field, indexer) not in where:
             # TODO should only be called once the graph is started
-            raise NoProviderException("Nobody provides {}: {}{}".format(kind, field, f"-{indexer}" if indexer else ""))
+            raise NoProviderException("Nobody provides {}: {}{}".format(kind, field, f"-{indexer}" if _has_indexer(indexer) else ""))
 
     def override(self, field: str, value: Any) -> None:
         raise NotImplementedError()

csp_gateway/server/gateway/csp/factory.py

@@ -103,7 +103,7 @@ def build(self, channels: ChannelsType) -> ChannelsType:
 
         # second pass to finish connecting in wires
         for (field, indexer), push_adapter in channels._send_channels.items():
-            with channels._connection_context(f"Send[{field}]{f'<{indexer}>' if indexer else ''}"):
+            with channels._connection_context(f"Send[{field}]{f'<{indexer}>' if indexer not in (None, '') else ''}"):
                 # Add it as an edge on the StreamGroup
                 if isinstance(push_adapter, GenericPushAdapter):
                     channels.set_channel(field, push_adapter.out(), indexer=indexer)

csp_gateway/server/shared/json_converter.py

@@ -93,7 +93,7 @@ def _create_snapshot_dict(all_data: list[ChannelValueModel]) -> dict[str, Any]:
     for data in all_data:
         channel = data.channel
         value = _convert_orjson_compatible(data.value)
-        if key := data.dict_basket_key:
+        if (key := data.dict_basket_key) not in (None, ""):
             res[channel][_convert_orjson_compatible(key)] = value
         else:
             res[channel] = value

csp_gateway/testing/harness.py

@@ -191,7 +191,12 @@ class GatewayAssertAttrUnsetEvent(BaseGatewayTestEvent):
     attr: str
 
     def apply(self, now, values, tick_counts, *args, **kwargs):
-        assert hasattr(values[self.channel], self.attr) == (not self.unset)
+        value = values[self.channel]
+        if hasattr(value, "_implicitly_unset"):
+            is_set = self.attr not in value._implicitly_unset() and getattr(value, self.attr, None) is not None
+        else:
+            is_set = hasattr(value, self.attr)
+        assert is_set == (not self.unset)
 
 
 class GatewayAssertLenEvent(BaseGatewayTestEvent):
@@ -256,7 +261,12 @@ class GatewayAssertIdxAttrUnsetEvent(BaseGatewayTestEvent):
     attr: str
 
     def apply(self, now, values, tick_counts, *args, **kwargs):
-        assert hasattr(values[self.channel][self.idx], self.attr) == (not self.unset)
+        value = values[self.channel][self.idx]
+        if hasattr(value, "_implicitly_unset"):
+            is_set = self.attr not in value._implicitly_unset() and getattr(value, self.attr, None) is not None
+        else:
+            is_set = hasattr(value, self.attr)
+        assert is_set == (not self.unset)
 
 
 class GatewayAssertTickedEvents(BaseGatewayTestEvent):

csp_gateway/tests/server/gateway/csp/test_channels.py

@@ -66,6 +66,12 @@ class MyGatewayChannels(GatewayChannels):
     my_array_channel: ts[Numpy1DArray[float]] = None
 
 
+def test_csp_enum_dict_basket_has_static_keys():
+    channels = MyGatewayChannels()
+
+    assert set(channels._keys_for_channel("my_enum_basket")) == set(MyEnum)
+
+
 class DerivedChannels(MyGatewayChannels):
     pass
 

csp_gateway/tests/server/gateway/test_gateway.py

@@ -3,7 +3,7 @@
 import time
 from collections.abc import Callable
 from datetime import datetime, timedelta
-from enum import Enum
+from enum import IntEnum
 from io import StringIO
 from typing import Annotated, Any
 
@@ -29,10 +29,14 @@
 from csp_gateway.utils import NoProviderException
 
 
-class MyEnum(Enum):
+class MyEnum(IntEnum):
+    ZERO = 0
     ONE = 1
     TWO = 2
 
+    def __str__(self):
+        return f"{type(self).__name__}.{self.name}"
+
 
 class MyStruct(GatewayStruct):
     foo: float
@@ -89,16 +93,19 @@ def connect(self, channels: MyGatewayChannels) -> None:
         channels.set_channel(MyGatewayChannels.my_array_channel, csp.const(np.array([1.0, 2.0])))
 
         if self.by_key:
+            channels.set_channel(MyGatewayChannels.my_enum_basket, self.my_data, MyEnum.ZERO)
             channels.set_channel(MyGatewayChannels.my_enum_basket, self.my_data, MyEnum.ONE)
             channels.set_channel(MyGatewayChannels.my_enum_basket, self.my_data2, MyEnum.TWO)
             channels.set_channel(MyGatewayChannels.my_str_basket, self.my_data, "my_key")
             channels.set_channel(MyGatewayChannels.my_str_basket, self.my_data2, "my_key2")
+            channels.set_channel(MyGatewayChannels.my_enum_basket_list, self.my_list_data, MyEnum.ZERO)
             channels.set_channel(MyGatewayChannels.my_enum_basket_list, self.my_list_data, MyEnum.ONE)
             channels.set_channel(MyGatewayChannels.my_enum_basket_list, self.my_list_data, MyEnum.TWO)
         else:
             channels.set_channel(
                 MyGatewayChannels.my_enum_basket,
                 {
+                    MyEnum.ZERO: self.my_data,
                     MyEnum.ONE: self.my_data,
                     MyEnum.TWO: self.my_data2,
                 },
@@ -113,6 +120,7 @@ def connect(self, channels: MyGatewayChannels) -> None:
             channels.set_channel(
                 MyGatewayChannels.my_enum_basket_list,
                 {
+                    MyEnum.ZERO: self.my_list_data,
                     MyEnum.ONE: self.my_list_data,
                     MyEnum.TWO: self.my_list_data,
                 },
@@ -185,6 +193,10 @@ def connect(self, channels: MyGatewayChannels) -> None:
             csp.add_graph_output(f"my_enum_basket_list[{k}]", v)
 
         # Get by indexer
+        csp.add_graph_output(
+            "my_enum_basket_ZERO",
+            channels.get_channel(MyGatewayChannels.my_enum_basket, MyEnum.ZERO),
+        )
         csp.add_graph_output(
             "my_enum_basket_ONE",
             channels.get_channel(MyGatewayChannels.my_enum_basket, MyEnum.ONE),
@@ -585,7 +597,9 @@ def test_last(by_key):
 
         output = gateway.channels.last("my_enum_basket")
         assert isinstance(output, dict)
-        assert len(output) == 2
+        assert len(output) == 3
+        output = gateway.channels.last("my_enum_basket", MyEnum.ZERO)
+        assert isinstance(output, MyStruct)
         output = gateway.channels.last("my_enum_basket", MyEnum.ONE)
         assert isinstance(output, MyStruct)
 
@@ -598,7 +612,7 @@ def test_last(by_key):
 
         output = gateway.channels.last("my_enum_basket_list")
         assert isinstance(output, dict)
-        assert len(output) == 2
+        assert len(output) == 3
     finally:
         gateway.stop()
 

csp_gateway/tests/server/modules/io/test_json_converter.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timedelta, timezone
+from enum import IntEnum
 from typing import Annotated, Any
 
 import csp
@@ -290,6 +291,17 @@ def test_convert_orjson_compatible():
     assert _convert_orjson_compatible(my_str) == my_str
 
 
+def test_create_snapshot_dict_with_zero_valued_enum_key():
+    class ZeroBasedEnum(IntEnum):
+        ZERO = 0
+
+    value = MyStruct(foo=1.0)
+    snapshot_dict = _create_snapshot_dict(
+        [CVM(channel="enum_basket", value=value, dict_basket_key=ZeroBasedEnum.ZERO, timestamp=datetime(2020, 1, 1))]
+    )
+    assert snapshot_dict["enum_basket"] == {"ZERO": _convert_orjson_compatible(value)}
+
+
 def test_parse_snapshot_dict():
     timestamp = datetime(2020, 1, 1)
     dummy_id = "9"

csp_gateway/tests/server/web/test_webserver.py

@@ -283,7 +283,7 @@ def test_csp_last_basket(self, rest_client: TestClient):
         assert len(data) == 3
 
         for channel_name, member in ExampleEnum.__members__.items():
-            multiplier = member.value
+            multiplier = getattr(member, "value", member)
             response = rest_client.get(f"/api/v1/last/basket/{channel_name}?token=test")
             assert response.status_code == 200
 
@@ -345,7 +345,7 @@ def test_csp_next_basket(self, rest_client: TestClient):
         assert len(data) == 3
 
         for channel_name, member in ExampleEnum.__members__.items():
-            multiplier = member.value
+            multiplier = getattr(member, "value", member)
             response = rest_client.get(f"/api/v1/next/basket/{channel_name}?token=test")
             assert response.status_code == 200
 

csp_gateway/tests/test_harness.py

@@ -8,6 +8,16 @@
 from csp_gateway.testing.shared_helpful_classes import MyGateway, MyGatewayChannels, MyStruct
 
 
+def test_assert_attr_unset_uses_gateway_struct_field_presence():
+    h = GatewayTestHarness(test_channels=[MyGatewayChannels.my_channel])
+    h.send(MyGatewayChannels.my_channel, MyStruct())
+    h.assert_attr_unset(MyGatewayChannels.my_channel, "foo")
+    h.assert_attr_unset(MyGatewayChannels.my_channel, "my_flag", unset=False)
+
+    gateway = MyGateway(modules=[h], channels=MyGatewayChannels())
+    csp.run(gateway.graph, starttime=datetime(2020, 1, 1), endtime=timedelta(0))
+
+
 @pytest.mark.parametrize("make_invalid", (True, False))
 def test_delay(make_invalid):
     channels = [