Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/py/mat3ra/code/entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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)
Expand Down
15 changes: 15 additions & 0 deletions tests/py/unit/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
from enum import Enum

from mat3ra.code.entity import InMemoryEntityPydantic
from pydantic import BaseModel
Expand Down Expand Up @@ -53,3 +54,17 @@ class ExampleNestedKeyAsClassInstanceClass(ExampleNestedSchema, InMemoryEntityPy

class ExampleDoubleNestedKeyAsClassInstancesClass(ExampleDoubleNestedSchema, InMemoryEntityPydantic):
double_nested_key1: ExampleNestedKeyAsClassInstanceClass


class SampleEnum(str, Enum):
VALUE1 = "value1"
VALUE2 = "value2"


class SampleModelWithEnum(BaseModel):
type: SampleEnum
name: str = "test"


class SampleEntityWithEnum(SampleModelWithEnum, InMemoryEntityPydantic):
pass
12 changes: 12 additions & 0 deletions tests/py/unit/test_entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
ExampleNestedKeyAsClassInstanceClass,
ExampleNestedSchema,
ExampleSchema,
SampleEnum,
SampleEntityWithEnum,
)


Expand Down Expand Up @@ -165,6 +167,16 @@ def test_to_dict():
assert result_exclude == {"key1": "value1"}


def test_to_dict_with_enum():
entity = SampleEntityWithEnum(type=SampleEnum.VALUE1, name="example")
result = entity.to_dict()

assert isinstance(result, dict)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

assert not instance

assert not isinstance(result["type"], SampleEnum) # Should not be an enum object
assert result == {"type": "value1", "name": "example"}
assert result["type"] == "value1" # String, not enum object


def test_to_json():
entity = ExampleClass.create(REFERENCE_OBJECT_VALID)
result = entity.to_json()
Expand Down
Loading