Add support for encoding schema as tag:

ketozhang · Nov 15, 2024 · 0d3a4f1 · 0d3a4f1
1 parent a810bb2
commit 0d3a4f1
Showing 1 changed file with 10 additions and 3 deletions.
diff --git 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})