Support CSP IntEnum compatibility

ccflow/enums.py

@@ -69,7 +69,10 @@ def validate(cls, v) -> "Enum":
         if isinstance(v, cls):
             return v
         elif isinstance(v, str):
-            return cls[v]
+            try:
+                return cls[v]
+            except KeyError as e:
+                raise ValueError(f"Cannot convert value to enum: {v}") from e
         elif isinstance(v, int):
             return cls(v)
         raise ValueError(f"Cannot convert value to enum: {v}")

ccflow/serialization.py

@@ -3,10 +3,10 @@
 import numpy as np
 import orjson
 
-from .enums import Enum
+from .enums import _CSP_ENUM, Enum
 
 
-def _remove_dict_enums(obj: Any) -> dict:
+def _remove_dict_enums(obj: Any) -> Any:
     if isinstance(obj, Enum):
         return obj.name
     elif isinstance(obj, dict):
@@ -21,21 +21,14 @@ def orjson_dumps(v, default=None, *arga, **kwargs) -> str:
     # with the json_encoders as the first argument. We try to perform the
     # conversion
     options = orjson.OPT_NON_STR_KEYS | orjson.OPT_NAIVE_UTC | orjson.OPT_SERIALIZE_NUMPY
-    try:
-        return orjson.dumps(
-            v,
-            default=default,
-            option=options,
-        ).decode()
-    except orjson.JSONEncodeError:
-        # if we fail, we try to remove the enums because
-        # orjson serialization fails when csp enums are
-        # used as dict keys. See https://github.com/ijl/orjson/issues/445
-        return orjson.dumps(
-            _remove_dict_enums(v),
-            default=default,
-            option=options,
-        ).decode()
+    # orjson serialization fails when legacy csp enums are used as dict keys, while IntEnum-based csp enums
+    # serialize as integers. Convert both representations to names first. See https://github.com/ijl/orjson/issues/445
+    value = _remove_dict_enums(v) if _CSP_ENUM else v
+    return orjson.dumps(
+        value,
+        default=default,
+        option=options,
+    ).decode()
 
 
 def make_ndarray_orjson_valid(arr: np.ndarray) -> list[Any] | np.ndarray:

ccflow/tests/enums/test_enums.py

@@ -96,6 +96,8 @@ class MyEnum(Enum):
         self.assertTrue(issubclass(MyEnum, Enum))
         self.assertTrue(isinstance(MyEnum.A.name, str))
         self.assertEqual(MyEnum[MyEnum.A.name], MyEnum.A)
+        with self.assertRaises(ValueError):
+            MyEnum.validate("missing")
 
         class MyAutoEnum(Enum):
             A = auto()

ccflow/tests/test_base_serialize.py

@@ -11,10 +11,10 @@
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict, Field, ValidationError
 
 from ccflow import BaseModel, GenericResult, NDArray
-from ccflow.enums import Enum
+from ccflow.enums import _CSP_ENUM, Enum
 from ccflow.exttypes.pydantic_numpy.ndtypes import bool_, complex64, float32, float64, int8, uint32
 from ccflow.pickling import reduce_generic_model_instance
-from ccflow.serialization import make_ndarray_orjson_valid
+from ccflow.serialization import make_ndarray_orjson_valid, orjson_dumps
 
 
 class ParentModel(BaseModel):
@@ -168,6 +168,10 @@ def test_serialization_nested(self):
     def test_serialization_enum(self):
         self._check_serialization(D(value=MyEnum.FIRST))
 
+    def test_orjson_serialization_enum_dict(self):
+        expected = '{"FIRST":"SECOND"}' if _CSP_ENUM else '{"1":2}'
+        self.assertEqual(orjson_dumps({MyEnum.FIRST: MyEnum.SECOND}), expected)
+
     def test_serialization_nested_subclass(self):
         self._check_serialization(NestedModel(a=ChildModel(field1=0, field2=10)))