Add support for self-references in JSON Schema

Author: Fatal1ty

mashumaro/jsonschema/models.py

@@ -212,3 +212,5 @@ class Context:
     plugins: Sequence[BasePlugin] = ()
     # PEP 695 TypeAliasType recursion guard
     _building_type_aliases: set[int] = field(default_factory=set, repr=False)
+    # Dataclass recursion guard (e.g. typing.Self, direct or mutual references)
+    _building_dataclasses: set[type] = field(default_factory=set, repr=False)

mashumaro/jsonschema/schema.py

@@ -41,6 +41,7 @@
     is_not_required,
     is_readonly,
     is_required,
+    is_self,
     is_special_typing_primitive,
     is_type_alias_type,
     is_type_var,
@@ -364,52 +365,69 @@ def override_with_any(reason: Any) -> None:
 
 @register
 def on_dataclass(instance: Instance, ctx: Context) -> Optional[JSONSchema]:
-    # TODO: Self references might not work
     if is_dataclass(instance.origin_type):
+        # When dataclasses reference themselves (typing.Self) or each other,
+        # we must break infinite recursion by forcing $ref/$defs.
+        origin = instance.origin_type
+
         if ctx.all_refs:
-            title = clean_id(type_name(instance.type, short=True))
-            title = title.strip("_")
+            def_key = clean_id(type_name(instance.type, short=True)).strip("_")
         else:
-            title = instance.origin_type.__name__
-        jsonschema_config = instance.get_self_config().json_schema
-        schema = JSONObjectSchema(
-            title=title,
-            additionalProperties=jsonschema_config.get(
-                "additionalProperties", False
-            ),
-        )
-        properties: dict[str, JSONSchema] = {}
-        required = []
-        field_schema_overrides = jsonschema_config.get("properties", {})
-        for f_name, f_type, has_default, f_default in instance.fields():
-            override = field_schema_overrides.get(f_name)
-            f_instance = instance.derive(type=f_type, name=f_name)
-            if override:
-                f_schema = JSONSchema.from_dict(override)
+            def_key = origin.__name__
+
+        ref_prefix = ctx.ref_prefix or ctx.dialect.definitions_root_pointer
+
+        if origin in ctx._building_dataclasses:
+            # Ensure placeholder exists so the final schema can fill it in.
+            ctx.definitions.setdefault(def_key, EmptyJSONSchema())
+            return JSONSchema(reference=f"{ref_prefix}/{def_key}")
+
+        ctx._building_dataclasses.add(origin)
+        try:
+            # If a placeholder exists (recursion), we'll populate it later
+            jsonschema_config = instance.get_self_config().json_schema
+            schema = JSONObjectSchema(
+                title=def_key,
+                additionalProperties=jsonschema_config.get(
+                    "additionalProperties", False
+                ),
+            )
+            properties: dict[str, JSONSchema] = {}
+            required = []
+            field_schema_overrides = jsonschema_config.get("properties", {})
+            for f_name, f_type, has_default, f_default in instance.fields():
+                override = field_schema_overrides.get(f_name)
+                f_instance = instance.derive(type=f_type, name=f_name)
+                if override:
+                    f_schema = JSONSchema.from_dict(override)
+                else:
+                    f_schema = get_schema(f_instance, ctx)
+                if f_instance.alias:
+                    f_name = f_instance.alias
+                if f_default is not MISSING:
+                    f_schema.default = f_default
+                description = f_instance.metadata.get("description")
+                if description:
+                    f_schema.description = description
+
+                if not has_default:
+                    required.append(f_name)
+
+                properties[f_name] = f_schema
+            if properties:
+                schema.properties = properties
+            if required:
+                schema.required = required
+
+            # If recursion was detected, we need $defs/$ref regardless
+            existing = ctx.definitions.get(def_key)
+            if ctx.all_refs or isinstance(existing, EmptyJSONSchema):
+                ctx.definitions[def_key] = schema
+                return JSONSchema(reference=f"{ref_prefix}/{def_key}")
             else:
-                f_schema = get_schema(f_instance, ctx)
-            if f_instance.alias:
-                f_name = f_instance.alias
-            if f_default is not MISSING:
-                f_schema.default = f_default
-            description = f_instance.metadata.get("description")
-            if description:
-                f_schema.description = description
-
-            if not has_default:
-                required.append(f_name)
-
-            properties[f_name] = f_schema
-        if properties:
-            schema.properties = properties
-        if required:
-            schema.required = required
-        if ctx.all_refs:
-            ctx.definitions[title] = schema
-            ref_prefix = ctx.ref_prefix or ctx.dialect.definitions_root_pointer
-            return JSONSchema(reference=f"{ref_prefix}/{title}")
-        else:
-            return schema
+                return schema
+        finally:
+            ctx._building_dataclasses.discard(origin)
 
 
 @register
@@ -466,8 +484,16 @@ def on_special_typing_primitive(
         )
     elif is_literal(instance.type):
         return on_literal(instance, ctx)
-    # elif is_self(instance.type):
-    #     raise NotImplementedError
+    elif is_self(instance.type):
+        # typing.Self / typing_extensions.Self is only meaningful inside
+        # a class body. In dataclasses, Instance.owner_class points to the
+        # dataclass that defines the field.
+        owner = instance.owner_class
+        if owner is None:  # pragma: no cover
+            raise NotImplementedError(
+                "typing.Self is supported only for dataclass fields"
+            )
+        return get_schema(instance.derive(type=owner), ctx)
     elif is_required(instance.type) or is_not_required(instance.type):
         return get_schema(instance.derive(type=args[0]), ctx)
     elif is_unpack(instance.type):

tests/test_jsonschema/test_self.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from typing_extensions import Self
+
+from mashumaro.jsonschema import build_json_schema
+from mashumaro.jsonschema.models import JSONArraySchema, JSONObjectSchema
+
+
+def test_jsonschema_supports_self_with_refs() -> None:
+    @dataclass
+    class Node:
+        child: Optional[Self] = None
+        items: list[Self] = None  # type: ignore[assignment]
+
+    schema = build_json_schema(Node, all_refs=True)
+
+    # Top-level schema should be a $ref into $defs
+    assert schema.reference is not None
+    assert schema.reference.startswith("#/$defs/")
+
+    def_key = schema.reference.split("#/$defs/", 1)[1]
+
+    # And the referenced definition should exist in $defs
+    assert schema.definitions is not None
+    assert def_key in schema.definitions
+
+    node_def = schema.definitions[def_key]
+    assert isinstance(node_def, JSONObjectSchema)
+
+    # Optional[Self] => anyOf [$ref-to-self, null]
+    child_schema = node_def.properties["child"]
+    assert child_schema.anyOf is not None
+    assert any(s.reference == schema.reference for s in child_schema.anyOf)
+    assert any(
+        getattr(s.type, "value", None) == "null" for s in child_schema.anyOf
+    )
+
+    # list[Self] => array items $ref-to-self
+    items_schema = node_def.properties["items"]
+    assert isinstance(items_schema, JSONArraySchema)
+    assert items_schema.items is not None
+    assert items_schema.items.reference == schema.reference
+
+
+def test_jsonschema_self_forces_refs_when_recursive_even_without_all_refs() -> (
+    None
+):
+    @dataclass
+    class Node:
+        child: Self | None = None
+
+    schema = build_json_schema(Node)
+
+    # For recursive dataclasses we expect $defs/$ref even if all_refs=False
+    assert schema.reference == "#/$defs/Node"
+    assert schema.definitions is not None
+    assert "Node" in schema.definitions
+
+    node_def = schema.definitions["Node"].to_dict()
+    assert node_def.get("title") == "Node"
+    assert node_def.get("type") == "object"
+    assert node_def.get("properties", {}).get("child", {}).get("anyOf")