From 4970e72914ada493ce19700bfaa570799539803a Mon Sep 17 00:00:00 2001
From: "Keto D. Zhang" <keto.zhang@gmail.com>
Date: Tue, 8 Oct 2024 21:49:44 -0700
Subject: [PATCH] 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()