diff --git a/ccflow/enums.py b/ccflow/enums.py index 44003b2..c12e4f1 100644 --- a/ccflow/enums.py +++ b/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}") diff --git a/ccflow/serialization.py b/ccflow/serialization.py index deb79b3..5f7e959 100644 --- a/ccflow/serialization.py +++ b/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: diff --git a/ccflow/tests/enums/test_enums.py b/ccflow/tests/enums/test_enums.py index 82604c2..f799ece 100644 --- a/ccflow/tests/enums/test_enums.py +++ b/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() diff --git a/ccflow/tests/test_base_serialize.py b/ccflow/tests/test_base_serialize.py index 3321632..797cc3a 100644 --- a/ccflow/tests/test_base_serialize.py +++ b/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)))