update: generate classes

Author: VsevolodX

src/py/mat3ra/code/entity.py

@@ -33,7 +33,10 @@ def __init__(self, code: ValidationErrorCode, details: Optional[ErrorDetails] =
 
 
 class InMemoryEntityPydantic(BaseModel):
-    model_config = {"arbitrary_types_allowed": True}
+    model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
+
+    # Factory helper field mapping field names to class names
+    _class_factory: Dict = {}
 
     @classmethod
     def create(cls: Type[T], config: Dict[str, Any]) -> T:
@@ -61,6 +64,18 @@ def clean(cls: Type[T], config: Dict[str, Any]) -> Dict[str, Any]:
         validated_model = cls.model_validate(config)
         return validated_model.model_dump()
 
+    def model_post_init(self, __context: Any) -> None:
+        for field_name, field_value in self.__dict__.items():
+            if isinstance(field_value, BaseModel):
+                class_reference = self._class_factory.get(field_name)
+                if not class_reference:
+                    continue
+                else:
+                    instance_field_name = field_name + "_instance"
+                    config = field_value.model_dump()  # convert from BaseModel to dict
+                    class_instance = class_reference(**config)
+                    setattr(self, instance_field_name, class_instance)
+
     def get_cls_name(self) -> str:
         return self.__class__.__name__
 

tests/py/unit/test_entity.py

@@ -6,17 +6,26 @@
 REFERENCE_OBJECT_VALID = {"key1": "value1", "key2": 1}
 REFERENCE_OBJECT_INVALID = {"key1": "value1", "key2": "value2"}
 REFERENCE_OBJECT_VALID_JSON = json.dumps(REFERENCE_OBJECT_VALID)
+REFERENCE_OBJECT_NESTED_VALID = {"nested_key1": {**REFERENCE_OBJECT_VALID}}
 
 
 class ExampleSchema(BaseModel):
     key1: str
     key2: int
 
 
+class ExampleNestedSchema(BaseModel):
+    nested_key1: ExampleSchema
+
+
 class ExampleClass(ExampleSchema, InMemoryEntityPydantic):
     pass
 
 
+class ExampleNestedClass(ExampleNestedSchema, InMemoryEntityPydantic):
+    _class_factory = {"nested_key1": ExampleClass}
+
+
 example_class_instance_valid = ExampleClass(**REFERENCE_OBJECT_VALID)
 
 
@@ -27,6 +36,16 @@ def test_create():
     assert in_memory_entity.key2 == 1
 
 
+def test_create_nested():
+    # Test creating an instance with nested valid data
+    in_memory_entity = ExampleNestedClass.create(REFERENCE_OBJECT_NESTED_VALID)
+    assert isinstance(in_memory_entity, ExampleNestedClass)
+    assert isinstance(in_memory_entity.nested_key1, ExampleSchema)
+    assert in_memory_entity.nested_key1.key1 == "value1"
+    assert in_memory_entity.nested_key1.key2 == 1
+    assert isinstance(in_memory_entity.nested_key1_instance, ExampleClass)
+
+
 def test_validate():
     # Test valid case
     in_memory_entity = ExampleClass.create(REFERENCE_OBJECT_VALID)