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()