feat: Add custom ASDF schema generator (#23)

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)

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

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,17 +25,19 @@ 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")
             )
         }
     )
     asdf.get_config().add_extension(TestExtension())
 
 
 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"

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"]