From 481c80428982dd557aeec8acf28059eefce2ec67 Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Sun, 6 Oct 2024 20:59:09 -0700 Subject: [PATCH 1/6] test: add test for astropy tables --- tests/examples/test_astropy_tables.py | 80 +++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 tests/examples/test_astropy_tables.py diff --git a/tests/examples/test_astropy_tables.py b/tests/examples/test_astropy_tables.py new file mode 100644 index 0000000..b6e498f --- /dev/null +++ b/tests/examples/test_astropy_tables.py @@ -0,0 +1,80 @@ +from __future__ import annotations +from typing import Annotated, TypeVar + +import asdf +import astropy.units as u +from asdf.extension import Extension +from astropy.table import Table +from astropy.units import Quantity +from pydantic import WithJsonSchema +import pytest +import yaml + +from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel + +T = TypeVar("T") + +AsdfAstropyTable = Annotated[ + T, + WithJsonSchema( + { + "type": "object", + "$ref": "http://asdf.org/table/table-1.1.0", + } + ), +] + + +class Database(AsdfPydanticModel): + _tag = "asdf://asdf-pydantic/examples/tags/database-1.0.0" + positions: AsdfAstropyTable[Table] + + +@pytest.fixture() +def asdf_extension(): + """Registers an ASDF extension containing models for this test.""" + AsdfPydanticConverter.add_models(Database) + + class TestExtension(Extension): + extension_uri = "asdf://asdf-pydantic/examples/extensions/test-1.0.0" + + converters = [AsdfPydanticConverter()] # type: ignore + tags = [*AsdfPydanticConverter().tags] # type: ignore + + asdf.get_config().add_extension(TestExtension()) + + with asdf.config_context() as asdf_config: + asdf_config.add_resource_mapping( + { + yaml.safe_load(Database.model_asdf_schema())[ + "id" + ]: Database.model_asdf_schema() + } + ) + print(Database.model_asdf_schema()) + asdf_config.add_extension(TestExtension()) + yield asdf_config + + +@pytest.mark.usefixtures("asdf_extension") +def test_convert_to_asdf(tmp_path): + database = Database( + positions=Table( + { + "x": Quantity([1, 2, 3], u.m), + "y": Quantity([4, 5, 6], u.m), + } + ) + ) + asdf.AsdfFile({"data": database}).write_to(tmp_path / "test.asdf") + + with asdf.open(tmp_path / "test.asdf") as af: + assert isinstance(af.tree["data"], Database) + assert isinstance(af.tree["data"].positions, Table) + + +@pytest.mark.usefixtures("asdf_extension") +def test_check_schema(): + """Tests the model schema is correct.""" + schema = yaml.safe_load(Database.model_asdf_schema()) + asdf.schema.check_schema(schema) From 166f4ec9350c85159d1bff1b405e98a5e6ee46b7 Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Sun, 6 Oct 2024 21:16:17 -0700 Subject: [PATCH 2/6] fix: tag url prefix --- tests/examples/test_astropy_tables.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/examples/test_astropy_tables.py b/tests/examples/test_astropy_tables.py index b6e498f..2d0998c 100644 --- a/tests/examples/test_astropy_tables.py +++ b/tests/examples/test_astropy_tables.py @@ -12,14 +12,14 @@ from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel -T = TypeVar("T") +T = TypeVar("T", bound=Table) AsdfAstropyTable = Annotated[ T, WithJsonSchema( { "type": "object", - "$ref": "http://asdf.org/table/table-1.1.0", + "$ref": "http://stsci.edu/schemas/asdf.org/table/table-1.1.0", } ), ] From 4970e72914ada493ce19700bfaa570799539803a Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Tue, 8 Oct 2024 21:49:44 -0700 Subject: [PATCH 3/6] feat: add WithAsdfSchema and AsdfTag for defining field's ASDF schema using annotations --- asdf_pydantic/schema.py | 41 +++++++++++++++++++++++++++ tests/examples/test_astropy_tables.py | 25 ++++++---------- 2 files changed, 49 insertions(+), 17 deletions(-) diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py index 38a2821..a870203 100644 --- a/asdf_pydantic/schema.py +++ b/asdf_pydantic/schema.py @@ -1,5 +1,36 @@ +""" + +## Adding existing ASDF tags as a field +Type annotation must be added to the field to specify the ASDF tag to use in the +ASDF schema. There are a few options to do this: + + - Use `AsdfTag` to specify the tag URI. + - Use `WithAsdfSchema` and pass in a dictionary to extend the schema with + additional properties. The key `"$ref"` can be used to specify the tag URI. + + from asdf_pydantic import AsdfPydanticModel + from asdf_pydantic.schema import AsdfTag + from astropy.table import Table + + class MyModel(AsdfPydanticModel): + table: Annotated[Table, AsdfTag("http://stsci.edu/schemas/asdf.org/table/table-1.1.0")] + +For more customization of the ASDF schema output, you can use `WithAsdfSchema` to +extend the schema with additional properties. + + # Changing the title of the field + table: Annotated[ + Table, + WithAsdfSchema({ + "title": "TABLE", + "$ref": "http://stsci.edu/schemas/asdf.org/table/table-1.1.0" + }), + ] +""" + from typing import Optional +from pydantic import WithJsonSchema from pydantic.json_schema import GenerateJsonSchema DEFAULT_ASDF_SCHEMA_REF_TEMPLATE = "#/definitions/{model}" @@ -60,3 +91,13 @@ def generate(self, schema, mode="validation"): } return json_schema + + +class WithAsdfSchema(WithJsonSchema): + def __init__(self, asdf_schema: dict, **kwargs): + json_schema = {"type": "object", **asdf_schema} + super().__init__(json_schema, **kwargs) + + +def AsdfTag(tag: str) -> WithAsdfSchema: + return WithAsdfSchema({"$ref": tag}) diff --git a/tests/examples/test_astropy_tables.py b/tests/examples/test_astropy_tables.py index 2d0998c..d90bd4a 100644 --- a/tests/examples/test_astropy_tables.py +++ b/tests/examples/test_astropy_tables.py @@ -1,33 +1,24 @@ from __future__ import annotations -from typing import Annotated, TypeVar + +from typing import Annotated import asdf import astropy.units as u +import pytest +import yaml from asdf.extension import Extension from astropy.table import Table from astropy.units import Quantity -from pydantic import WithJsonSchema -import pytest -import yaml from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel - -T = TypeVar("T", bound=Table) - -AsdfAstropyTable = Annotated[ - T, - WithJsonSchema( - { - "type": "object", - "$ref": "http://stsci.edu/schemas/asdf.org/table/table-1.1.0", - } - ), -] +from asdf_pydantic.schema import AsdfTag class Database(AsdfPydanticModel): _tag = "asdf://asdf-pydantic/examples/tags/database-1.0.0" - positions: AsdfAstropyTable[Table] + positions: Annotated[ + Table, AsdfTag("http://stsci.edu/schemas/asdf.org/table/table-1.1.0") + ] @pytest.fixture() From c84652f5013e8c6f1c378b5d8a5b54bf924c3aa3 Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Fri, 15 Nov 2024 11:09:29 -0800 Subject: [PATCH 4/6] feat!: migrate example models to tests/ and remove example extension from package fixes #37 --- asdf_pydantic/examples/__init__.py | 0 asdf_pydantic/examples/extensions.py | 19 ----------- pyproject.toml | 3 -- tests/convert_to_asdf_yaml_tree_test.py | 2 +- tests/entry_point_test.py | 11 ------- {asdf_pydantic => tests}/examples/shapes.py | 0 tests/examples/test_rectangle.py | 9 ++---- {asdf_pydantic => tests}/examples/tree.py | 0 tests/schema_validation_test.py | 35 +++++++++++++++++++-- 9 files changed, 36 insertions(+), 43 deletions(-) delete mode 100644 asdf_pydantic/examples/__init__.py delete mode 100644 asdf_pydantic/examples/extensions.py delete mode 100644 tests/entry_point_test.py rename {asdf_pydantic => tests}/examples/shapes.py (100%) rename {asdf_pydantic => tests}/examples/tree.py (100%) diff --git a/asdf_pydantic/examples/__init__.py b/asdf_pydantic/examples/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/asdf_pydantic/examples/extensions.py b/asdf_pydantic/examples/extensions.py deleted file mode 100644 index b79d4d9..0000000 --- a/asdf_pydantic/examples/extensions.py +++ /dev/null @@ -1,19 +0,0 @@ -from __future__ import annotations - -from asdf.extension import Extension - -from asdf_pydantic.converter import AsdfPydanticConverter -from asdf_pydantic.examples.shapes import AsdfRectangle -from asdf_pydantic.examples.tree import AsdfTreeNode - -AsdfPydanticConverter.add_models(AsdfRectangle, AsdfTreeNode) - - -class ExampleExtension(Extension): - extension_uri = "asdf://asdf-pydantic/examples/extensions/examples-1.0.0" - converters = [AsdfPydanticConverter()] # type: ignore - tags = [*AsdfPydanticConverter().tags] # type: ignore - - -def get_extensions() -> list[Extension]: - return [ExampleExtension()] diff --git a/pyproject.toml b/pyproject.toml index a5393e3..34148d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,9 +33,6 @@ Documentation = "https://asdf-pydantic.readthedocs.io" Issues = "https://github.com/ketozhang/asdf-pydantic/issues" Source = "https://github.com/ketozhang/asdf-pydantic" -[project.entry-points] -'asdf.extensions' = { asdf_pydantic_extension = 'asdf_pydantic.examples.extensions:get_extensions' } - [tool.hatch.version] source = "vcs" diff --git a/tests/convert_to_asdf_yaml_tree_test.py b/tests/convert_to_asdf_yaml_tree_test.py index 37825e1..2f9510f 100644 --- a/tests/convert_to_asdf_yaml_tree_test.py +++ b/tests/convert_to_asdf_yaml_tree_test.py @@ -1,6 +1,6 @@ from __future__ import annotations -from asdf_pydantic.examples.tree import AsdfTreeNode, Node +from tests.examples.tree import AsdfTreeNode, Node def test_sanity(): diff --git a/tests/entry_point_test.py b/tests/entry_point_test.py deleted file mode 100644 index efccb39..0000000 --- a/tests/entry_point_test.py +++ /dev/null @@ -1,11 +0,0 @@ -from tempfile import NamedTemporaryFile - -import asdf - -from asdf_pydantic.examples.shapes import AsdfRectangle - - -def test_create_asdf_file(): - with NamedTemporaryFile() as tempfile: - af = asdf.AsdfFile({"rect": AsdfRectangle(width=42, height=10)}) - af.write_to(tempfile.name) diff --git a/asdf_pydantic/examples/shapes.py b/tests/examples/shapes.py similarity index 100% rename from asdf_pydantic/examples/shapes.py rename to tests/examples/shapes.py diff --git a/tests/examples/test_rectangle.py b/tests/examples/test_rectangle.py index a95a2ba..9d414d7 100644 --- a/tests/examples/test_rectangle.py +++ b/tests/examples/test_rectangle.py @@ -7,13 +7,8 @@ import yaml from asdf.extension import Extension -from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel - - -class AsdfRectangle(AsdfPydanticModel): - _tag = "asdf://asdf-pydantic/examples/tags/rectangle-1.0.0" - width: float - height: float +from asdf_pydantic import AsdfPydanticConverter +from tests.examples.shapes import AsdfRectangle @pytest.fixture() diff --git a/asdf_pydantic/examples/tree.py b/tests/examples/tree.py similarity index 100% rename from asdf_pydantic/examples/tree.py rename to tests/examples/tree.py diff --git a/tests/schema_validation_test.py b/tests/schema_validation_test.py index a5365b5..60a3f09 100644 --- a/tests/schema_validation_test.py +++ b/tests/schema_validation_test.py @@ -1,4 +1,5 @@ from tempfile import NamedTemporaryFile +from typing import Annotated import asdf import pydantic @@ -7,8 +8,9 @@ from asdf.extension import Extension from asdf_pydantic import AsdfPydanticConverter -from asdf_pydantic.examples.shapes import AsdfRectangle -from asdf_pydantic.examples.tree import AsdfTreeNode +from asdf_pydantic.model import AsdfPydanticModel +from tests.examples.shapes import AsdfRectangle +from tests.examples.tree import AsdfTreeNode def setup_module(): @@ -136,3 +138,32 @@ def test_given_child_field_contains_asdf_object_then_schema_has_child_tag(): child_schema = schema["definitions"]["AsdfNode"]["properties"]["child"] assert {"tag": AsdfTreeNode._tag} in child_schema["anyOf"] + + +######################################################################################## +# AsdfTag +######################################################################################## +from asdf_pydantic.schema import AsdfTag # noqa: E402 + + +@pytest.mark.parametrize( + "asdf_tag_str, mode, expected_ref_key", + [ + ("http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", "auto", "$ref"), + ("http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", "ref", "$ref"), + ("tag:stsci.edu:asdf/table/table-1.1.0", "auto", "tag"), + ("tag:stsci.edu:asdf/table/table-1.1.0", "tag", "tag"), + ], +) +def test_tag_mode(asdf_tag_str: str, mode, expected_ref_key): + """Test that schema correctly has ``$ref:`` or ``tag:`` depending on the + selected mode. + """ + from astropy.table import Table + + class TestModel(AsdfPydanticModel): + _tag = "asdf://asdf-pydantic/examples/tags/test-model-1.0.0" + table: Annotated[Table, AsdfTag(asdf_tag_str, mode=mode)] + + schema = yaml.safe_load(TestModel.model_asdf_schema()) + assert expected_ref_key in schema["properties"]["table"] From 738514c37a3ec634feb4a6eedc86bd48b4087fd9 Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Fri, 15 Nov 2024 11:48:38 -0800 Subject: [PATCH 5/6] fix: remove type field from schema if ref exists --- asdf_pydantic/schema.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py index a870203..838e0e1 100644 --- a/asdf_pydantic/schema.py +++ b/asdf_pydantic/schema.py @@ -95,8 +95,8 @@ def generate(self, schema, mode="validation"): class WithAsdfSchema(WithJsonSchema): def __init__(self, asdf_schema: dict, **kwargs): - json_schema = {"type": "object", **asdf_schema} - super().__init__(json_schema, **kwargs) + super().__init__(asdf_schema, **kwargs) + def AsdfTag(tag: str) -> WithAsdfSchema: From f0fe466bfe8ae22039f8687d1562613a4747447d Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Fri, 15 Nov 2024 11:50:00 -0800 Subject: [PATCH 6/6] Add support for encoding schema as `tag:` --- asdf_pydantic/schema.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py index 838e0e1..705cbe7 100644 --- a/asdf_pydantic/schema.py +++ b/asdf_pydantic/schema.py @@ -28,7 +28,7 @@ class MyModel(AsdfPydanticModel): ] """ -from typing import Optional +from typing import Literal, Optional from pydantic import WithJsonSchema from pydantic.json_schema import GenerateJsonSchema @@ -98,6 +98,13 @@ def __init__(self, asdf_schema: dict, **kwargs): super().__init__(asdf_schema, **kwargs) +def AsdfTag(tag: str, mode: Literal["auto", "ref", "tag"] = "auto") -> WithAsdfSchema: + if mode == "auto": + parsed_mode = "tag" if tag.startswith("tag") else "ref" + else: + parsed_mode = mode -def AsdfTag(tag: str) -> WithAsdfSchema: - return WithAsdfSchema({"$ref": tag}) + if parsed_mode == "tag": + return WithAsdfSchema({"tag": tag}) + else: + return WithAsdfSchema({"$ref": tag})