diff --git a/.gitignore b/.gitignore index 6a73c32..311267d 100644 --- a/.gitignore +++ b/.gitignore @@ -135,3 +135,6 @@ dmypy.json # Pyre type checker .pyre/ + +# uv +uv.lock diff --git a/asdf_pydantic/converter.py b/asdf_pydantic/converter.py index 4be2684..d81383f 100644 --- a/asdf_pydantic/converter.py +++ b/asdf_pydantic/converter.py @@ -40,4 +40,4 @@ def to_yaml_tree(self, obj: AsdfPydanticModel, tag, ctx): return obj.asdf_yaml_tree() def from_yaml_tree(self, node, tag, ctx): - return self._tag_to_class[tag].parse_obj(node) + return self._tag_to_class[tag].model_validate(node) diff --git a/asdf_pydantic/model.py b/asdf_pydantic/model.py index 2d38157..20e2918 100644 --- a/asdf_pydantic/model.py +++ b/asdf_pydantic/model.py @@ -15,33 +15,42 @@ class AsdfPydanticModel(BaseModel): ASDF Serialization and Deserialization: Serialize to ASDF yaml tree is done with the py:classmethod`AsdfPydanticModel.asdf_yaml_tree()` and deserialize to an - AsdfPydanticModel object with py:meth`AsdfPydanticModel.parse_obj()`. + AsdfPydanticModel object with py:meth`AsdfPydanticModel.model_validate()`. """ _tag: ClassVar[str | TagDefinition] model_config = ConfigDict(arbitrary_types_allowed=True) def asdf_yaml_tree(self) -> dict: - d = {} - for field_key, v in self.__dict__.items(): - if field_key not in self.__fields__: - continue - + """Converts the model to an ASDF-compatible YAML tree (dict). + + .. note:: + Any fields that are normal Pydantic `BaseModel` will be converted to + dict. See conversion table. + + Conversion Table: + + +-------------------------+-----------------------------+ + | Value type in field | Value type in dict | + +=========================+=============================+ + | AsdfPydanticModel | No conversion | + +-------------------------+-----------------------------+ + | BaseModel | Converted to dict using | + | | BaseModel.model_dump() | + +-------------------------+-----------------------------+ + | Other types | No conversion | + +-------------------------+-----------------------------+ + """ + tree = {} + for k, v in dict(self).items(): if isinstance(v, AsdfPydanticModel): - d[field_key] = v + tree[k] = v + elif isinstance(v, BaseModel): + tree[k] = v.model_dump() else: - d[field_key] = self._get_value( - v, - to_dict=True, - by_alias=False, - include=None, - exclude=None, - exclude_unset=False, - exclude_defaults=False, - exclude_none=False, - ) - - return d + tree[k] = v + + return tree @model_validator(mode="before") @classmethod diff --git a/tests/test_model.py b/tests/test_model.py index 9f78227..00f1f3d 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1,10 +1,52 @@ import pytest import yaml from asdf.extension import TagDefinition +from pydantic import BaseModel from asdf_pydantic import AsdfPydanticModel +######################################################################################## +# ASDF Tree +######################################################################################## +def test_asdf_yaml_tree_sanity(): + class TestModel(AsdfPydanticModel): + foo: str + + model = TestModel(foo="bar") + assert model.asdf_yaml_tree() == {"foo": "bar"} + + +def test_asdf_yaml_tree_nested_dict(): + class NestedTestModel(AsdfPydanticModel): + nested: dict + + model = NestedTestModel(nested={"foo": "bar"}) + assert model.asdf_yaml_tree() == {"nested": {"foo": "bar"}} + + +def test_asdf_yaml_tree_nested_basemodel(): + class TestModel(BaseModel): + foo: str + + class NestedTestModel(AsdfPydanticModel): + nested: TestModel + + model = NestedTestModel(nested=TestModel(foo="bar")) + assert model.asdf_yaml_tree() == {"nested": {"foo": "bar"}} + + +def test_asdf_yaml_tree_nested_asdf_pydantic_model(): + class TestModel(AsdfPydanticModel): + foo: str + + class NestedTestModel(AsdfPydanticModel): + nested: TestModel + + model = NestedTestModel(nested=TestModel(foo="bar")) + assert model.asdf_yaml_tree() == {"nested": TestModel(foo="bar")} + + ######################################################################################## # TAGS ########################################################################################