diff --git a/asdf_pydantic/model.py b/asdf_pydantic/model.py index 0bdc49a..7de635d 100644 --- a/asdf_pydantic/model.py +++ b/asdf_pydantic/model.py @@ -1,8 +1,10 @@ -import textwrap from typing import ClassVar import yaml from pydantic import BaseModel +from typing_extensions import deprecated + +from asdf_pydantic.schema import DEFAULT_ASDF_SCHEMA_REF_TEMPLATE, GenerateAsdfSchema class AsdfPydanticModel(BaseModel): @@ -42,29 +44,45 @@ def asdf_yaml_tree(self) -> dict: return d @classmethod + def model_asdf_schema( + cls, + by_alias: bool = True, + ref_template: str = DEFAULT_ASDF_SCHEMA_REF_TEMPLATE, + schema_generator: type[GenerateAsdfSchema] = GenerateAsdfSchema, + ): + """Get the ASDF schema definition for this model.""" + # Implementation follows closely with the `BaseModel.model_json_schema` + schema_generator_instance = schema_generator( + by_alias=by_alias, + ref_template=ref_template, + tag=cls._tag, + ) + json_schema = schema_generator_instance.generate(cls.__pydantic_core_schema__) + + header = "%YAML 1.1\n---\n" + + return f"{header}\n{yaml.safe_dump(json_schema)}" + + @classmethod + @deprecated( + "The `schema_asdf` method is deprecated; use `model_asdf_schema` instead." + ) def schema_asdf( - cls, *, metaschema: str = "http://stsci.edu/schemas/asdf/asdf-schema-1.0.0" + cls, + *, + metaschema: str = GenerateAsdfSchema.schema_dialect, + **kwargs, ) -> str: """Get the ASDF schema definition for this model. Parameters ---------- metaschema, optional - A metaschema URI, by default "http://stsci.edu/schemas/asdf/asdf-schema-1.0.0". - See https://asdf.readthedocs.io/en/stable/asdf/extending/schemas.html#anatomy-of-a-schema - for more options. + A metaschema URI """ # noqa: E501 - # TODO: Function signature should follow BaseModel.schema() or - # BaseModel.schema_json() - header = textwrap.dedent( - f""" - %YAML 1.1 - --- - $schema: {metaschema} - id: {cls._tag} - tag: tag:{cls._tag.split('://', maxsplit=2)[-1]} - - """ - ) - body = yaml.safe_dump(cls.model_json_schema()) - return header + body + if metaschema != GenerateAsdfSchema.schema_dialect: + raise NotImplementedError( + f"Only {GenerateAsdfSchema.schema_dialect} is supported as metaschema." + ) + + return cls.model_asdf_schema(**kwargs) diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py new file mode 100644 index 0000000..321144f --- /dev/null +++ b/asdf_pydantic/schema.py @@ -0,0 +1,40 @@ +from typing import Optional + +from pydantic.json_schema import GenerateJsonSchema + +DEFAULT_ASDF_SCHEMA_REF_TEMPLATE = "#/definitions/{model}" + + +class GenerateAsdfSchema(GenerateJsonSchema): + """Generates ASDF-compatible schema from Pydantic's default JSON schema generator. + + ```{caution} Experimental + This schema generator is not complete. It currently creates JSON 2020-12 + schema (despite `$schema` says it's `asdf-schema-1.0.0`) which are not + compatible with ASDF. + ``` + """ + + # HACK: When we can support tree models, then not all schema should have tag + tag: Optional[str] + schema_dialect = "http://stsci.edu/schemas/asdf/asdf-schema-1.0.0" + + def __init__( + self, + by_alias: bool = True, + ref_template: str = DEFAULT_ASDF_SCHEMA_REF_TEMPLATE, + tag: Optional[str] = None, + ): + super().__init__(by_alias=by_alias, ref_template=ref_template) + self.tag = tag + + def generate(self, schema, mode="validation"): + json_schema = super().generate(schema, mode) # noqa: F841 + + if self.tag: + json_schema["$schema"] = self.schema_dialect + json_schema["id"] = self.tag + json_schema["tag"] = f"tag:{self.tag.split('://', maxsplit=2)[-1]}" + + # TODO: Convert jsonschema 2020-12 to ASDF schema + return json_schema diff --git a/tests/examples/test_rectangle.py b/tests/examples/test_rectangle.py index b52ea49..bd24368 100644 --- a/tests/examples/test_rectangle.py +++ b/tests/examples/test_rectangle.py @@ -1,6 +1,8 @@ import asdf +import pytest from asdf.extension import Extension from asdf.schema import check_schema, load_schema +from yaml.scanner import ScannerError from asdf_pydantic import AsdfPydanticConverter from asdf_pydantic.examples.shapes import AsdfRectangle @@ -23,7 +25,7 @@ class TestExtension(Extension): asdf.get_config().add_resource_mapping( { "asdf://asdf-pydantic/shapes/schemas/rectangle-1.0.0": ( - AsdfRectangle.schema_asdf().encode("utf-8") + AsdfRectangle.model_asdf_schema().encode("utf-8") ) } ) @@ -31,9 +33,11 @@ class TestExtension(Extension): def test_schema(): - schema = load_schema("asdf://asdf-pydantic/shapes/schemas/rectangle-1.0.0") - - check_schema(schema) + try: + schema = load_schema("asdf://asdf-pydantic/shapes/schemas/rectangle-1.0.0") + check_schema(schema) + except ScannerError as e: + pytest.fail(f"{e}\n{AsdfRectangle.model_asdf_schema()}") assert schema["$schema"] == "http://stsci.edu/schemas/asdf/asdf-schema-1.0.0" assert schema["title"] == "AsdfRectangle" diff --git a/tests/schema_validation_test.py b/tests/schema_validation_test.py index f62ac0d..85cc05e 100644 --- a/tests/schema_validation_test.py +++ b/tests/schema_validation_test.py @@ -28,7 +28,7 @@ class TestExtension(Extension): asdf.get_config().add_resource_mapping( { "asdf://asdf-pydantic/shapes/schemas/rectangle-1.0.0": ( - AsdfRectangle.schema_asdf().encode("utf-8") + AsdfRectangle.model_asdf_schema().encode("utf-8") ) } ) @@ -130,7 +130,7 @@ def test_validate_fail_on_bad_yaml_file(): def test_given_child_field_contains_asdf_object_then_schema_has_child_tag(): from asdf.schema import check_schema - schema = yaml.safe_load(AsdfNode.schema_asdf()) # type: ignore + schema = yaml.safe_load(AsdfNode.model_asdf_schema()) # type: ignore check_schema(schema) child_schema = schema["definitions"]["AsdfNode"]["properties"]["child"]