update: fix serialization of enums + tests

Author: VsevolodX

src/py/mat3ra/code/entity.py

@@ -62,7 +62,7 @@ def from_json(cls: Type[T], json_str: str) -> T:
     def clean(cls: Type[T], config: Dict[str, Any]) -> Dict[str, Any]:
         # Validate the config; extra keys are dropped and defaults are substituted.
         validated = cls.model_validate(config, strict=False)
-        return validated.model_dump(exclude_unset=False)
+        return validated.model_dump(mode="json", exclude_unset=False)
 
     def get_schema(self) -> Dict[str, Any]:
         return self.model_json_schema()
@@ -77,7 +77,7 @@ def get_cls_name(self) -> str:
         return self.__class__.__name__
 
     def to_dict(self, exclude: Optional[List[str]] = None) -> Dict[str, Any]:
-        return self.model_dump(exclude=set(exclude) if exclude else None)
+        return self.model_dump(mode="json", exclude=set(exclude) if exclude else None)
 
     def to_json(self, exclude: Optional[List[str]] = None) -> str:
         return self.model_dump_json(exclude=set(exclude) if exclude else None)

tests/py/unit/__init__.py

@@ -1,4 +1,5 @@
 import json
+from enum import Enum
 
 from mat3ra.code.entity import InMemoryEntityPydantic
 from pydantic import BaseModel
@@ -53,3 +54,17 @@ class ExampleNestedKeyAsClassInstanceClass(ExampleNestedSchema, InMemoryEntityPy
 
 class ExampleDoubleNestedKeyAsClassInstancesClass(ExampleDoubleNestedSchema, InMemoryEntityPydantic):
     double_nested_key1: ExampleNestedKeyAsClassInstanceClass
+
+
+class SampleEnum(str, Enum):
+    TRI = "TRI"
+    ANGSTROM = "angstrom"
+
+
+class SampleModelWithEnum(BaseModel):
+    type: SampleEnum
+    name: str = "test"
+
+
+class SampleEntityWithEnum(SampleModelWithEnum, InMemoryEntityPydantic):
+    pass

tests/py/unit/test_entity.py

@@ -18,6 +18,8 @@
     ExampleNestedKeyAsClassInstanceClass,
     ExampleNestedSchema,
     ExampleSchema,
+    SampleEnum,
+    SampleEntityWithEnum,
 )
 
 
@@ -165,6 +167,26 @@ def test_to_dict():
     assert result_exclude == {"key1": "value1"}
 
 
+def test_to_dict_with_enum():
+    """Test that enums are serialized as strings in to_dict()"""
+    entity = SampleEntityWithEnum(type=SampleEnum.TRI, name="example")
+    result = entity.to_dict()
+    
+    assert isinstance(result, dict)
+    assert result == {"type": "TRI", "name": "example"}
+    assert result["type"] == "TRI"  # String, not enum object
+
+
+def test_to_dict_with_enum_exclude():
+    """Test to_dict with exclude parameter works with enums"""
+    entity = SampleEntityWithEnum(type=SampleEnum.ANGSTROM, name="test")
+    result = entity.to_dict(exclude=["name"])
+    
+    assert isinstance(result, dict)
+    assert result == {"type": "angstrom"}
+    assert result["type"] == "angstrom"  # String, not enum object
+
+
 def test_to_json():
     entity = ExampleClass.create(REFERENCE_OBJECT_VALID)
     result = entity.to_json()