Merge pull request #303 from Fatal1ty/recursive-type-jsonschema

Add support for recursive types in JSON Schema

Author: Fatal1ty

mashumaro/jsonschema/models.py

@@ -210,3 +210,5 @@ class Context:
     all_refs: Optional[bool] = None
     ref_prefix: Optional[str] = None
     plugins: Sequence[BasePlugin] = ()
+    # PEP 695 TypeAliasType recursion guard
+    _building_type_aliases: set[int] = field(default_factory=set, repr=False)

mashumaro/jsonschema/schema.py

@@ -317,6 +317,15 @@ class Config(config_cls):  # type: ignore
 register = Registry.register
 
 
+def _type_alias_definition_name(alias_type: Any) -> str:
+    """Return a stable $defs key for PEP 695 TypeAliasType."""
+
+    name = getattr(alias_type, "__name__", None)
+    if isinstance(name, str) and name:
+        return name
+    return clean_id(str(id(alias_type)))
+
+
 BASIC_TYPES = {str, int, float, bool}
 
 
@@ -474,7 +483,36 @@ def on_special_typing_primitive(
         if evaluated is not None:
             return get_schema(instance.derive(type=evaluated), ctx)
     elif is_type_alias_type(instance.type):
-        return get_schema(instance.derive(type=instance.type.__value__), ctx)
+        alias_type = instance.type
+        def_name = _type_alias_definition_name(alias_type)
+        alias_id = id(alias_type)
+        ref_prefix = ctx.ref_prefix or ctx.dialect.definitions_root_pointer
+
+        # If we're already building this alias, it's recursion.
+        # In that case, force using $ref/$defs.
+        if alias_id in ctx._building_type_aliases:
+            # The $defs placeholder may not exist yet (mutual recursion).
+            ctx.definitions.setdefault(def_name, EmptyJSONSchema())
+            return JSONSchema(reference=f"{ref_prefix}/{def_name}")
+
+        ctx._building_type_aliases.add(alias_id)
+        try:
+            value_schema = get_schema(
+                instance.derive(type=alias_type.__value__), ctx
+            )
+        finally:
+            ctx._building_type_aliases.discard(alias_id)
+
+        # If the alias is marked as recursive (direct or mutual),
+        # store its definition in $defs and return a $ref.
+        if def_name in ctx.definitions or ctx.all_refs:
+            existing = ctx.definitions.get(def_name)
+            if existing is None or isinstance(existing, EmptyJSONSchema):
+                ctx.definitions[def_name] = value_schema
+            return JSONSchema(reference=f"{ref_prefix}/{def_name}")
+
+        # Non-recursive alias: return the schema directly, without $defs.
+        return value_schema
 
 
 @register

tests/test_jsonschema/test_jsonschema_pep_695.py

@@ -1,9 +1,137 @@
-type MyTypeAliasType = int | str
+from mashumaro.core.meta.types.common import clean_id
 from mashumaro.jsonschema import build_json_schema
+from mashumaro.jsonschema.models import Context
+from mashumaro.jsonschema.schema import _type_alias_definition_name
+
+type JSON = int | str | float | bool | None | list[JSON] | dict[str, JSON]
+type X = int | str
+type A = int | list[B]
+type B = str | list[A]
 
 
 def test_type_alias_type_with_jsonschema():
-    schema = build_json_schema(MyTypeAliasType)
+    schema = build_json_schema(X)
     assert schema.to_dict() == {
         "anyOf": [{"type": "integer"}, {"type": "string"}]
     }
+
+
+def test_jsonschema_for_recursive_union() -> None:
+    schema = build_json_schema(JSON)
+    assert schema.to_dict() == {
+        "$ref": "#/$defs/JSON",
+        "$defs": {
+            "JSON": {
+                "anyOf": [
+                    {"type": "integer"},
+                    {"type": "string"},
+                    {"type": "number"},
+                    {"type": "boolean"},
+                    {"type": "null"},
+                    {"type": "array", "items": {"$ref": "#/$defs/JSON"}},
+                    {
+                        "type": "object",
+                        "additionalProperties": {"$ref": "#/$defs/JSON"},
+                        "propertyNames": {"type": "string"},
+                    },
+                ]
+            }
+        },
+    }
+
+
+def test_jsonschema_for_mutual_recursive_type_aliases_without_refs() -> None:
+    schema = build_json_schema(A)
+    assert schema.to_dict() == {
+        "$ref": "#/$defs/A",
+        "$defs": {
+            "A": {
+                "anyOf": [
+                    {"type": "integer"},
+                    {
+                        "type": "array",
+                        "items": {
+                            "anyOf": [
+                                {"type": "string"},
+                                {
+                                    "type": "array",
+                                    "items": {"$ref": "#/$defs/A"},
+                                },
+                            ]
+                        },
+                    },
+                ]
+            }
+        },
+    }
+
+
+def test_jsonschema_for_mutual_recursive_type_aliases_with_refs() -> None:
+    schema = build_json_schema(A, all_refs=True)
+    assert schema.to_dict() == {
+        "$ref": "#/$defs/A",
+        "$defs": {
+            "A": {
+                "anyOf": [
+                    {"type": "integer"},
+                    {"type": "array", "items": {"$ref": "#/$defs/B"}},
+                ]
+            },
+            "B": {
+                "anyOf": [
+                    {"type": "string"},
+                    {"type": "array", "items": {"$ref": "#/$defs/A"}},
+                ]
+            },
+        },
+    }
+
+
+def test_type_alias_non_recursive_inlines_when_all_refs_false() -> None:
+    schema = build_json_schema(X, all_refs=False)
+    assert schema.to_dict() == {
+        "anyOf": [
+            {"type": "integer"},
+            {"type": "string"},
+        ]
+    }
+
+
+def test_type_alias_non_recursive_uses_defs_when_all_refs_true() -> None:
+    schema = build_json_schema(X, all_refs=True)
+    assert schema.to_dict() == {
+        "$ref": "#/$defs/X",
+        "$defs": {
+            "X": {
+                "anyOf": [
+                    {"type": "integer"},
+                    {"type": "string"},
+                ]
+            }
+        },
+    }
+
+
+def test_type_alias_placeholder_not_leaking_into_context_defs() -> None:
+    # Ensure that for non-recursive aliases with all_refs=False we don't leave
+    # unused entries in Context.definitions.
+    ctx = Context()
+    schema = build_json_schema(X, context=ctx, all_refs=False)
+    assert schema.to_dict() == {
+        "anyOf": [
+            {"type": "integer"},
+            {"type": "string"},
+        ]
+    }
+    assert ctx.definitions == {}
+
+
+def test_type_alias_definition_name_falls_back_to_clean_id_when_name_empty() -> (
+    None
+):
+    class NamelessAlias:
+        __name__ = ""
+
+    alias = NamelessAlias()
+
+    assert _type_alias_definition_name(alias) == clean_id(str(id(alias)))