Support legacy CSP enum compatibility

csp_gateway/server/gateway/csp/factory.py

@@ -3,6 +3,7 @@
 from typing import Generic, get_args, get_origin
 
 from ccflow import BaseModel
+from csp import Enum as CspEnum
 from csp.impl.genericpushadapter import GenericPushAdapter
 from pydantic import Field
 
@@ -83,12 +84,13 @@ def build(self, channels: ChannelsType) -> ChannelsType:
 
         # Now wire in the signals
         # first pass is for any baskets
-        for (field, _indexer), push_adapter in channels._send_channels.items():
+        # Snapshot the items: add_send_channel below inserts the per-key entries into _send_channels.
+        for (field, _indexer), push_adapter in list(channels._send_channels.items()):
             if isinstance(push_adapter, tuple):
                 # dict basket, plug in now that we should know all the
                 # possible keys
                 key_type = get_dict_basket_key_type(channels.get_outer_type(field))
-                if isinstance(key_type, type) and issubclass(key_type, Enum):
+                if isinstance(key_type, type) and issubclass(key_type, (Enum, CspEnum)):
                     for enumfield in key_type:
                         channels.add_send_channel(field, enumfield)
 

csp_gateway/server/shared/json_converter.py

@@ -1,14 +1,14 @@
 import logging
 from collections import defaultdict, deque
 from datetime import UTC, datetime, timedelta
-from enum import Enum
+from enum import Enum as PyEnum
 from typing import Any, TypeVar
 
 import csp
 import numpy as np
 from ccflow import BaseModel
 from ccflow.serialization import make_ndarray_orjson_valid
-from csp import ts
+from csp import Enum as CspEnum, ts
 from csp.impl.types.tstype import isTsType
 from pydantic import BaseModel as PydanticBaseModel, Field, PrivateAttr, TypeAdapter
 
@@ -28,7 +28,7 @@
 
 logger = logging.getLogger(__name__)
 
-_KEY_TYPE = str | Enum
+_KEY_TYPE = str | PyEnum | CspEnum
 T = TypeVar("T")
 K = TypeVar("K")
 
@@ -78,7 +78,7 @@ def _convert_orjson_compatible(obj: Any):
         return {_convert_orjson_compatible(k): _convert_orjson_compatible(v) for k, v in obj.items()}
     if isinstance(obj, (list, set, tuple)):
         return [_convert_orjson_compatible(val) for val in obj]
-    if isinstance(obj, Enum):
+    if isinstance(obj, (PyEnum, CspEnum)):
         return obj.name
     if isinstance(obj, timedelta):
         return obj.total_seconds()

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

@@ -0,0 +1,62 @@
+from datetime import datetime, timedelta
+from enum import Enum as PyEnum
+
+import csp
+import pytest
+from csp import Enum as CspEnum, ts
+
+from csp_gateway import Gateway, GatewayChannels, GatewayModule, GatewayStruct
+
+
+class MyStruct(GatewayStruct):
+    foo: float
+
+
+class MyPyEnum(PyEnum):
+    A = 0
+    B = 1
+
+
+class MyCspEnum(CspEnum):
+    A = 0
+    B = 1
+
+
+class PyEnumChannels(GatewayChannels):
+    basket: dict[MyPyEnum, ts[MyStruct]] = None
+
+
+class CspEnumChannels(GatewayChannels):
+    basket: dict[MyCspEnum, ts[MyStruct]] = None
+
+
+class SendWholeBasketModule(GatewayModule):
+    key_type: type
+    seen: list = []
+
+    @csp.node
+    def _value(self, trigger: ts[bool]) -> ts[MyStruct]:
+        if csp.ticked(trigger):
+            return MyStruct(foo=1.0)
+
+    def connect(self, channels: GatewayChannels) -> None:
+        trigger = csp.timer(interval=timedelta(seconds=0.1), value=True)
+        for key in self.key_type:
+            channels.set_channel("basket", self._value(trigger), key)
+        # No indexer: the factory has to expand this into a send channel per key.
+        channels.add_send_channel("basket")
+        self.seen.append(channels)
+
+
+@pytest.mark.parametrize(
+    "channels_type,key_type",
+    [(PyEnumChannels, MyPyEnum), (CspEnumChannels, MyCspEnum)],
+    ids=["python_enum", "csp_enum"],
+)
+def test_whole_basket_send_channel_expands_enum_keys(channels_type, key_type):
+    module = SendWholeBasketModule(key_type=key_type, seen=[])
+    gateway = Gateway(modules=[module], channels=channels_type())
+    csp.run(gateway.graph, starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1))
+
+    channels = module.seen[0]
+    assert set(channels._send_channels) == {("basket", None)} | {("basket", key) for key in key_type}

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

@@ -24,6 +24,7 @@
     ChannelValueModel as CVM,
     _convert_orjson_compatible,
     _create_snapshot_dict,
+    _create_tuple_dict_basket,
 )
 from csp_gateway.testing.shared_helpful_classes import (
     MyEnum,
@@ -37,6 +38,11 @@
 class MysteryClass: ...
 
 
+class MyCspEnum(csp.Enum):
+    ZERO = 0
+    ONE = 1
+
+
 # This is an example on how to annotate a new, custom class, for validation and ser/der via pydantic
 # GatewayStructs will pick this up and allow users to use their own defined serializtion and deserialization
 ValidatedMysteryClass = Annotated[
@@ -287,6 +293,9 @@ def test_convert_orjson_compatible():
     enum = MyEnum.ONE
     assert _convert_orjson_compatible(enum) == enum.name
 
+    csp_enum = MyCspEnum.ONE
+    assert _convert_orjson_compatible(csp_enum) == csp_enum.name
+
     my_str = "hai"
     assert _convert_orjson_compatible(my_str) == my_str
 
@@ -302,6 +311,24 @@ class ZeroBasedEnum(IntEnum):
     assert snapshot_dict["enum_basket"] == {"ZERO": _convert_orjson_compatible(value)}
 
 
+def test_create_snapshot_dict_with_zero_valued_csp_enum_key():
+    value = MyStruct(foo=1.0)
+    snapshot_dict = _create_snapshot_dict([CVM(channel="enum_basket", value=value, dict_basket_key=MyCspEnum.ZERO, timestamp=datetime(2020, 1, 1))])
+    assert snapshot_dict["enum_basket"] == {"ZERO": _convert_orjson_compatible(value)}
+
+
+def test_create_tuple_dict_basket_with_csp_enum_key():
+    @csp.graph
+    def graph():
+        csp.add_graph_output(
+            "result",
+            _create_tuple_dict_basket("enum_basket", csp.const("value"), MyCspEnum.ZERO),
+        )
+
+    result = csp.run(graph, starttime=datetime(2020, 1, 1))
+    assert result["result"][0][1] == ("enum_basket", (MyCspEnum.ZERO, "value"))
+
+
 def test_parse_snapshot_dict():
     timestamp = datetime(2020, 1, 1)
     dummy_id = "9"

csp_gateway/tests/server/modules/web/test_perspective.py

@@ -1,6 +1,7 @@
 import json
 from datetime import UTC, date, datetime, timedelta
 from enum import Enum, auto
+from typing import Annotated
 from unittest.mock import MagicMock
 
 import csp
@@ -106,6 +107,40 @@ class ChildStruct(GatewayStruct, BaseModelWithContainer):
     assert schema["arr"] is int
 
 
+def test_csp_enum_schema():
+    class Status(csp.Enum):
+        READY = 1
+
+    class StructWithCspEnum(GatewayStruct):
+        status: Status
+
+    assert StructWithCspEnum.psp_schema()["status"] is str
+
+
+def test_annotated_csp_enum_array_schema():
+    class Status(csp.Enum):
+        READY = 1
+        DONE = 2
+
+    class StructWithCspEnumArray(GatewayStruct):
+        statuses: Annotated[Numpy1DArray[Status], Field(description="Statuses")]
+        values: Annotated[Numpy1DArray[float], Field(description="Values")]
+
+    schema = StructWithCspEnumArray.psp_schema()
+    assert schema["statuses"] is str
+    assert schema["values"] is float
+
+    value = StructWithCspEnumArray(statuses=np.array([Status.READY, Status.DONE], dtype=object))
+    assert [row["statuses"] for row in value.psp_flatten()] == ["READY", "DONE"]
+
+
+def test_optional_annotated_array_schema():
+    class StructWithOptionalArray(GatewayStruct):
+        values: Annotated[Numpy1DArray[float], Field(description="Values")] | None = None
+
+    assert StructWithOptionalArray.psp_schema()["values"] is float
+
+
 def test_inherited_ndarray_annotation_schema():
     # Base declares numpy array field with element type
     class BaseModelWithArray(BaseModel):

csp_gateway/tests/test_optional_imports.py

@@ -1,6 +1,14 @@
+import builtins
+import runpy
 import subprocess
 import sys
 import textwrap
+from enum import Enum
+from pathlib import Path
+
+import pytest
+
+FILTER_PATH = Path(__file__).parents[1] / "utils/web/filter.py"
 
 
 def _run_python(source: str) -> subprocess.CompletedProcess[str]:
@@ -65,3 +73,29 @@ def find_spec(self, fullname, path, target=None):
 
     assert result.returncode == 0, result.stderr
     assert result.stdout == "OK\n"
+
+
+def test_filter_allows_missing_csp(monkeypatch):
+    original_import = builtins.__import__
+
+    def missing_csp(name, *args, **kwargs):
+        if name == "csp":
+            raise ModuleNotFoundError("No module named 'csp'", name="csp")
+        return original_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", missing_csp)
+    namespace = runpy.run_path(FILTER_PATH)
+    assert namespace["_ENUM_TYPES"] == (Enum,)
+
+
+def test_filter_surfaces_broken_csp(monkeypatch):
+    original_import = builtins.__import__
+
+    def broken_csp(name, *args, **kwargs):
+        if name == "csp":
+            raise ImportError("csp internal dependency is broken", name="csp_internal")
+        return original_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", broken_csp)
+    with pytest.raises(ImportError, match="csp internal dependency is broken"):
+        runpy.run_path(FILTER_PATH)

csp_gateway/tests/utils/web/test_query.py

@@ -1,6 +1,8 @@
 import json
 from datetime import datetime
+from types import SimpleNamespace
 
+import csp
 import pytest
 
 from csp_gateway.server.demo import ExampleData
@@ -125,6 +127,14 @@ def test_query_simple(self):
         q = Query(filters=[Filter(attr="x", by=FilterCondition(value=0, where="=="))])
         assert q.calculate(DUMMY_STATE_DATA) == [d for d in DUMMY_STATE_DATA if d.x == 0]
 
+    def test_query_csp_enum_by_name(self):
+        class Status(csp.Enum):
+            READY = 1
+
+        item = SimpleNamespace(status=Status.READY)
+        q = Query(filters=[Filter(attr="status", by=FilterCondition(value="READY", where="=="))])
+        assert q.calculate([item]) == [item]
+
     def test_query_attr(self):
         q = Query(filters=[Filter(attr="id", by=FilterCondition(attr="y", where=">="))])
         assert q.calculate(DUMMY_STATE_DATA) == [d for d in DUMMY_STATE_DATA if d.id >= d.y]

csp_gateway/utils/struct/psp.py

@@ -8,6 +8,7 @@
 from typing import Annotated, Any, Optional, Union, get_args, get_origin
 
 import orjson
+from csp import Enum as CspEnum
 from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from numpy import ndarray
 from pydantic import BaseModel
@@ -237,6 +238,9 @@ def psp_schema(cls, excluded_columns: ExcludedColumns | None = None) -> dict[str
                         raise KeyError(field)
 
                 # get arg type
+                annotation = _strip_annotated(annotation)
+                if _is_optional(annotation):
+                    annotation = _get_type_from_optional(annotation)
                 arg = get_args(annotation)[0]
 
                 # use this as type
@@ -253,7 +257,7 @@ def psp_schema(cls, excluded_columns: ExcludedColumns | None = None) -> dict[str
             continue
 
         # If its an enum, promote to str
-        if issubclass(value, PyEnum):
+        if issubclass(value, (PyEnum, CspEnum)):
             schema[field] = str
             continue
 
@@ -316,7 +320,7 @@ def _callback(obj):
                 return obj.tolist()
             elif isinstance(obj, set):
                 return list(obj)
-            elif isinstance(obj, PyEnum):
+            elif isinstance(obj, (PyEnum, CspEnum)):
                 return obj.name
             elif isinstance(obj, _thread.LockType):
                 return "<Lock>"

csp_gateway/utils/web/filter.py

@@ -5,6 +5,16 @@
 
 from pydantic import BaseModel, Field
 
+# csp is a server-only dependency, but this module is imported by client-only installs.
+try:
+    from csp import Enum as CspEnum
+
+    _ENUM_TYPES = (PyEnum, CspEnum)
+except ModuleNotFoundError as error:
+    if error.name != "csp":
+        raise
+    _ENUM_TYPES = (PyEnum,)
+
 log = logging.getLogger(__name__)
 
 FilterWhere = Literal["==", "!=", "<", "<=", ">", ">="]
@@ -55,7 +65,7 @@ def calculate(self, obj) -> bool:
             if self.by.value is not None:
                 lhs = _get_nested_attr(obj, self.attr)
                 # Convert enums attrs to strings during filtering
-                if isinstance(lhs, PyEnum):
+                if isinstance(lhs, _ENUM_TYPES):
                     lhs = lhs.name
                 log.info(f"Filtering: {lhs} {self.by.where} {self.by.value}")
                 return FilterWhereLambdaMap[self.by.where](lhs, self.by.value)