From 343df24cd0d0dd761a401a8964d9f7eb25a4022b Mon Sep 17 00:00:00 2001 From: "Keto D. Zhang" Date: Fri, 15 Nov 2024 11:50:00 -0800 Subject: [PATCH] Add support for encoding schema as `tag:` --- asdf_pydantic/schema.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py index 838e0e1..705cbe7 100644 --- 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})