Allow enums and structs to be handled by pydantic (#428)

* Pydantic integration with csp Enums, Structs, and csp numpy array typing

---------

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

Author: NeejWeej

csp/impl/enum.py

@@ -64,6 +64,49 @@ def __repr__(self):
     def __str__(self):
         return f"{type(self).__name__}.{self.name}"
 
+    @classmethod
+    def _validate(cls, v) -> "Enum":
+        if isinstance(v, cls):
+            return v
+        elif isinstance(v, str):
+            return cls[v]
+        elif isinstance(v, int):
+            return cls(v)
+        raise ValueError(f"Cannot convert value to enum: {v}")
+
+    @staticmethod
+    def _serialize(value: typing.Union[str, "Enum"]) -> str:
+        return value.name
+
+    @classmethod
+    def __get_pydantic_json_schema__(cls, _core_schema, handler):
+        from pydantic_core import core_schema
+
+        field_schema = handler(core_schema.str_schema())
+        field_schema.update(
+            type="string",
+            title=cls.__name__,
+            description=cls.__doc__ or "An enumeration of {}".format(cls.__name__),
+            enum=list(cls.__members__.keys()),
+        )
+        return field_schema
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls,
+        source_type,
+        handler,
+    ):
+        from pydantic_core import core_schema
+
+        return core_schema.no_info_before_validator_function(
+            cls._validate,
+            core_schema.any_schema(),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls._serialize, info_arg=False, return_schema=core_schema.str_schema(), when_used="json"
+            ),
+        )
+
 
 def DynamicEnum(name: str, values: typing.Union[dict, list], start=0, module_name=None):
     """create a csp.Enum type dynamically

csp/impl/struct.py

@@ -68,7 +68,10 @@ def __new__(cls, name, bases, dct):
         dct["__metadata__"] = metadata
         dct["__defaults__"] = defaults
 
-        return super().__new__(cls, name, bases, dct)
+        res = super().__new__(cls, name, bases, dct)
+        # This is how we make sure we construct the pydantic schema from the new class
+        res.__get_pydantic_core_schema__ = classmethod(res._get_pydantic_core_schema)
+        return res
 
     def layout(self, num_cols=8):
         layout = super()._layout()
@@ -81,6 +84,66 @@ def layout(self, num_cols=8):
         out += layout[idx : idx + num_cols]
         return out
 
+    @staticmethod
+    def _get_pydantic_core_schema(cls, _source_type, handler):
+        """Tell Pydantic how to validate and serialize this Struct class."""
+        from pydantic import PydanticSchemaGenerationError
+        from pydantic_core import core_schema
+
+        fields = {}
+        for field_name, field_type in cls.__full_metadata_typed__.items():
+            if field_name.startswith("_"):
+                continue  # we skip fields with underscore, like pydantic does
+            try:
+                field_schema = handler.generate_schema(field_type)
+            except PydanticSchemaGenerationError:
+                # This logic allows for handling generic types with types we cant get a schema for, only 1 layer deep, same as csp
+                item_tp = typing.Any if typing.get_origin(field_type) is None else typing.get_args(field_type)[0]
+                try:
+                    field_schema = handler.generate_schema(item_tp)
+                except PydanticSchemaGenerationError:
+                    field_schema = core_schema.any_schema()  # give up finally
+
+            if field_name in cls.__defaults__:
+                field_schema = core_schema.with_default_schema(
+                    schema=field_schema, default=cls.__defaults__[field_name]
+                )
+            fields[field_name] = core_schema.typed_dict_field(
+                schema=field_schema,
+                required=False,  # Make all fields optional
+            )
+        # Schema for dictionary inputs
+        fields_schema = core_schema.typed_dict_schema(
+            fields=fields,
+            total=False,  # Allow missing fields
+            extra_behavior="allow",  # let csp catch extra attributes, allows underscore fields to pass through
+        )
+
+        def create_instance(raw_data, validator):
+            # We choose to not revalidate, this is the default behavior in pydantic
+            if isinstance(raw_data, cls):
+                return raw_data
+            try:
+                return cls(**validator(raw_data))
+            except AttributeError as e:
+                #  Pydantic can't use AttributeError to check other classes, like in Union annotations
+                raise ValueError(str(e)) from None
+
+        def serializer(val, handler):
+            # We don't use 'to_dict' since that works recursively, we ignore underscore leading fields
+            new_val = {
+                k: getattr(val, k) for k in val.__full_metadata_typed__ if not k.startswith("_") and hasattr(val, k)
+            }
+            return handler(new_val)
+
+        return core_schema.no_info_wrap_validator_function(
+            function=create_instance,
+            schema=fields_schema,
+            serialization=core_schema.wrap_serializer_function_ser_schema(
+                function=serializer, schema=fields_schema, when_used="always"
+            ),
+        )
+
 
 class Struct(_csptypesimpl.PyStruct, metaclass=StructMeta):
     @classmethod

csp/tests/impl/test_enum.py

@@ -1,6 +1,10 @@
 import _pickle
+import json
+import pytest
 import unittest
 from datetime import datetime, timedelta
+from pydantic import BaseModel, ConfigDict, RootModel
+from typing import Dict, List
 
 import csp
 from csp import ts
@@ -36,6 +40,22 @@ def s1():
 MyDEnum = csp.DynamicEnum("MyDEnum", ["A", "B", "C"])
 
 
+class MyEnum3(csp.Enum):
+    FIELD1 = csp.Enum.auto()
+    FIELD2 = csp.Enum.auto()
+
+
+class MyModel(BaseModel):
+    enum: MyEnum3
+    enum_default: MyEnum3 = MyEnum3.FIELD1
+
+
+class MyDictModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True)
+
+    enum_dict: Dict[MyEnum3, int] = None
+
+
 class TestCspEnum(unittest.TestCase):
     def test_basic(self):
         self.assertEqual(MyEnum("A"), MyEnum.A)
@@ -152,6 +172,73 @@ class B(A):
 
         self.assertEqual("Cannot extend csp.Enum 'A': inheriting from an Enum is prohibited", str(cm.exception))
 
+    def test_pydantic_validation(self):
+        assert MyModel(enum="FIELD2").enum == MyEnum3.FIELD2
+        assert MyModel(enum=0).enum == MyEnum3.FIELD1
+        assert MyModel(enum=MyEnum3.FIELD1).enum == MyEnum3.FIELD1
+        with pytest.raises(ValueError):
+            MyModel(enum=3.14)
+
+    def test_pydantic_dict(self):
+        assert dict(MyModel(enum=MyEnum3.FIELD2)) == {"enum": MyEnum3.FIELD2, "enum_default": MyEnum3.FIELD1}
+        assert MyModel(enum=MyEnum3.FIELD2).model_dump(mode="python") == {
+            "enum": MyEnum3.FIELD2,
+            "enum_default": MyEnum3.FIELD1,
+        }
+        assert MyModel(enum=MyEnum3.FIELD2).model_dump(mode="json") == {"enum": "FIELD2", "enum_default": "FIELD1"}
+
+    def test_pydantic_serialization(self):
+        assert "enum" in MyModel.model_fields
+        assert "enum_default" in MyModel.model_fields
+        tm = MyModel(enum=MyEnum3.FIELD2)
+        assert json.loads(tm.model_dump_json()) == json.loads('{"enum": "FIELD2", "enum_default": "FIELD1"}')
+
+    def test_enum_as_dict_key_json_serialization(self):
+        class DictWrapper(RootModel[Dict[MyEnum3, int]]):
+            model_config = ConfigDict(use_enum_values=True)
+
+            def __getitem__(self, item):
+                return self.root[item]
+
+        class MyDictWrapperModel(BaseModel):
+            model_config = ConfigDict(use_enum_values=True)
+
+            enum_dict: DictWrapper
+
+        dict_model = MyDictModel(enum_dict={MyEnum3.FIELD1: 8, MyEnum3.FIELD2: 19})
+        assert dict_model.enum_dict[MyEnum3.FIELD1] == 8
+        assert dict_model.enum_dict[MyEnum3.FIELD2] == 19
+
+        assert json.loads(dict_model.model_dump_json()) == json.loads('{"enum_dict":{"FIELD1":8,"FIELD2":19}}')
+
+        dict_wrapper_model = MyDictWrapperModel(enum_dict=DictWrapper({MyEnum3.FIELD1: 8, MyEnum3.FIELD2: 19}))
+
+        assert dict_wrapper_model.enum_dict[MyEnum3.FIELD1] == 8
+        assert dict_wrapper_model.enum_dict[MyEnum3.FIELD2] == 19
+        assert json.loads(dict_wrapper_model.model_dump_json()) == json.loads('{"enum_dict":{"FIELD1":8,"FIELD2":19}}')
+
+    def test_json_schema_csp(self):
+        assert MyModel.model_json_schema() == {
+            "properties": {
+                "enum": {
+                    "description": "An enumeration of MyEnum3",
+                    "enum": ["FIELD1", "FIELD2"],
+                    "title": "MyEnum3",
+                    "type": "string",
+                },
+                "enum_default": {
+                    "default": "FIELD1",
+                    "description": "An enumeration of MyEnum3",
+                    "enum": ["FIELD1", "FIELD2"],
+                    "title": "MyEnum3",
+                    "type": "string",
+                },
+            },
+            "required": ["enum"],
+            "title": "MyModel",
+            "type": "object",
+        }
+
 
 if __name__ == "__main__":
     unittest.main()