Reuse validated schema definitions (#3467)

Author: koxudaxi

src/datamodel_code_generator/parser/jsonschema.py

@@ -2467,7 +2467,7 @@ def _load_ref_schema_object(self, ref: str) -> JsonSchemaObject:
             return cached
 
         file_part, fragment = ([*resolved_ref.split("#", 1), ""])[:2]
-        raw_doc = self._get_ref_body(file_part) if file_part else self.raw_obj
+        raw_doc = self._get_ref_raw_doc(file_part)
 
         target_schema: dict[str, YamlValue] | YamlValue = raw_doc
         if fragment:
@@ -2478,6 +2478,15 @@ def _load_ref_schema_object(self, ref: str) -> JsonSchemaObject:
         self._ref_schema_object_cache[resolved_ref] = ref_schema
         return ref_schema
 
+    def _get_ref_raw_doc(self, file_part: str) -> dict[str, YamlValue]:
+        """Return the raw schema document for a resolved reference file part."""
+        match file_part:
+            case "":
+                return self.raw_obj
+            case current_root_ref if current_root_ref == "/".join(self.model_resolver.current_root):
+                return self.raw_obj
+        return self._get_ref_body(file_part)
+
     def _anchor_ref_path(self, root_key: tuple[str, ...], path: list[str]) -> str:  # noqa: PLR6301
         """Return the local ref path for an anchor under the current root."""
         root_len = len(root_key)
@@ -5987,6 +5996,7 @@ def _parse_file(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
                     seen_definition_metadata_paths.add(definition_path_key)
                     obj = self._validate_schema_object(model, definition_path)
                     validated_definition_objects[definition_path_key] = obj
+                    self._ref_schema_object_cache.setdefault(self.model_resolver.resolve_ref(definition_path), obj)
                     self.parse_id(obj, definition_path)
                     if obj.recursiveAnchor:
                         ref_path = self._anchor_ref_path(root_key, definition_path)

tests/data/expected/main/jsonschema/same_file_ref_fast_path.py

@@ -0,0 +1,15 @@
+# generated by datamodel-codegen:
+#   filename:  same_file_ref_fast_path.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class User(BaseModel):
+    name: str
+
+
+class SameFileRefFastPath(BaseModel):
+    user: User

tests/data/jsonschema/same_file_ref_fast_path.json

@@ -0,0 +1,25 @@
+{
+  "title": "SameFileRefFastPath",
+  "type": "object",
+  "properties": {
+    "user": {
+      "$ref": "#/definitions/User"
+    }
+  },
+  "required": [
+    "user"
+  ],
+  "definitions": {
+    "User": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "name"
+      ]
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2246,6 +2246,16 @@ def test_main_serialize_as_any_import_fast_path(output_file: Path) -> None:
     )
 
 
+def test_main_same_file_ref_fast_path(output_file: Path) -> None:
+    """Test same-file refs keep generated output unchanged."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "same_file_ref_fast_path.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+    )
+
+
 def test_main_root_model_with_additional_properties_literal(min_version: str, output_file: Path) -> None:
     """Test root model additional properties with literal types."""
     run_main_and_assert(

tests/parser/test_jsonschema.py

@@ -63,6 +63,20 @@ def _validate_schema_object(
         return super()._validate_schema_object(raw, path)
 
 
+class RefBodyCountingJsonSchemaParser(CountingJsonSchemaParser):
+    """JsonSchemaParser that records reference body loads."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize the parser and reference body recorder."""
+        super().__init__(*args, **kwargs)
+        self.ref_body_paths: list[str] = []
+        self.ref_bodies: dict[str, dict[str, Any]] = {}
+
+    def _get_ref_body(self, resolved_ref: str) -> dict[str, Any]:
+        self.ref_body_paths.append(resolved_ref)
+        return self.ref_bodies[resolved_ref]
+
+
 @pytest.fixture(autouse=True)
 def block_dns_by_default(mocker: MockerFixture) -> None:
     """Keep tests that mock httpx.get independent from external DNS."""
@@ -108,24 +122,88 @@ def test_load_ref_schema_object_caches_refs() -> None:
     ]
 
 
+def test_load_ref_schema_object_reuses_current_raw_obj_for_current_root_ref() -> None:
+    """Test same-document refs use the current raw object instead of loading a ref body."""
+    parser = RefBodyCountingJsonSchemaParser("")
+    parser.raw_obj = {
+        "definitions": {
+            "User": {
+                "type": "object",
+                "properties": {"name": {"type": "string"}},
+            },
+        },
+    }
+
+    with parser.model_resolver.current_root_context(["current.json"]):
+        schema = parser._load_ref_schema_object("#/definitions/User")
+
+    assert schema.type == "object"
+    assert parser.ref_body_paths == []
+    assert parser.validation_paths == [
+        ("current.json#/definitions/User",),
+    ]
+
+
+def test_load_ref_schema_object_loads_non_current_root_ref_body() -> None:
+    """Test refs outside the current root still load their referenced body."""
+    parser = RefBodyCountingJsonSchemaParser("")
+    parser.raw_obj = {}
+    parser.ref_bodies["other.json"] = {
+        "definitions": {
+            "User": {
+                "type": "object",
+                "properties": {"name": {"type": "string"}},
+            },
+        },
+    }
+
+    with parser.model_resolver.current_root_context(["current.json"]):
+        schema = parser._load_ref_schema_object("other.json#/definitions/User")
+
+    assert schema.type == "object"
+    assert parser.ref_body_paths == ["other.json"]
+    assert parser.validation_paths == [
+        ("other.json#/definitions/User",),
+    ]
+
+
 def test_parse_file_reuses_validated_definition_objects() -> None:
-    """Test definitions validated for metadata are reused when parsed."""
+    """Test definitions validated for metadata are reused when parsed and referenced."""
     parser = CountingJsonSchemaParser("")
     raw = {
+        "type": "object",
+        "properties": {
+            "user": {"$ref": "#/definitions/User"},
+            "member": {"$ref": "#/definitions/Namespaces/$defs/Member"},
+        },
         "definitions": {
             "User": {
                 "type": "object",
                 "properties": {"name": {"type": "string"}},
             },
+            "Namespaces": {
+                "$defs": {
+                    "Member": {
+                        "type": "object",
+                        "properties": {"id": {"type": "integer"}},
+                    },
+                },
+            },
         },
     }
 
     parser._parse_file(raw, "Root", [])
 
     assert parser.validation_paths.count(("#/definitions", "User")) == 1
+    assert parser.validation_paths.count(("#/definitions", "Namespaces", "$defs", "Member")) == 1
+    assert ("#/definitions/User",) not in parser.validation_paths
+    assert ("#/definitions/Namespaces/$defs/Member",) not in parser.validation_paths
     reference = parser.model_resolver.get(["#/definitions", "User"])
     assert reference is not None
     assert reference.loaded
+    nested_reference = parser.model_resolver.get(["#/definitions", "Namespaces", "$defs", "Member"])
+    assert nested_reference is not None
+    assert nested_reference.loaded
 
 
 @pytest.mark.parametrize(