Cache JSON Schema ref facts (#3472)

* Cache JSON Schema ref facts

* Harden ref path cache test

Author: koxudaxi

src/datamodel_code_generator/parser/jsonschema.py

@@ -903,6 +903,7 @@ def __init__(
         self._dynamic_anchor_index: dict[tuple[str, ...], dict[str, str]] = {}
         self._recursive_anchor_index: dict[tuple[str, ...], list[str]] = {}
         self._ref_data_type_facts: dict[str, tuple[Any, bool]] = {}
+        self._local_ref_path_cache: dict[Path, Path] = {}
         self._force_base_model_refs: set[str] = set()
         self._force_base_model_generation = False
         self.field_keys: set[str] = {
@@ -1994,6 +1995,12 @@ def _get_x_python_import_path(self, x_python_import: dict[str, Any]) -> str | No
             raise Error(msg)
         return _validate_schema_python_import_path(f"{module}.{type_name}", "x-python-import")
 
+    def _cache_ref_data_type_facts(self, resolved_ref: str, obj: JsonSchemaObject) -> None:
+        self._ref_data_type_facts[resolved_ref] = (
+            obj.extras.get("x-python-import"),
+            obj.type == "null" or (self.strict_nullable and obj.nullable is True),
+        )
+
     def get_ref_data_type(self, ref: str) -> DataType:
         """Get a data type from a reference string.
 
@@ -2461,18 +2468,20 @@ def _deep_merge(self, dict1: dict[Any, Any], dict2: dict[Any, Any]) -> dict[Any,
             result[key] = value
         return result
 
-    def _load_ref_schema_object(self, ref: str) -> JsonSchemaObject:
-        """Load a JsonSchemaObject from a $ref using standard resolve/load pipeline."""
-        resolved_ref = self.model_resolver.resolve_ref(ref)
+    def _get_ref_raw_schema(self, resolved_ref: str) -> dict[str, YamlValue] | YamlValue:
         file_part, fragment = ([*resolved_ref.split("#", 1), ""])[:2]
         raw_doc = self._get_ref_body(file_part) if file_part else self.raw_obj
 
         target_schema: dict[str, YamlValue] | YamlValue = raw_doc
         if fragment:
             pointer = split_json_pointer(raw_doc, fragment)
             target_schema = get_model_by_path(raw_doc, pointer)
+        return target_schema
 
-        return self._validate_schema_object(target_schema, [resolved_ref])
+    def _load_ref_schema_object(self, ref: str) -> JsonSchemaObject:
+        """Load a JsonSchemaObject from a $ref using standard resolve/load pipeline."""
+        resolved_ref = self.model_resolver.resolve_ref(ref)
+        return self._validate_schema_object(self._get_ref_raw_schema(resolved_ref), [resolved_ref])
 
     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."""
@@ -5181,9 +5190,13 @@ def _get_ref_body(self, resolved_ref: str) -> dict[str, YamlValue]:
         return self._get_ref_body_from_remote(resolved_ref)
 
     def _resolve_local_ref_path(self, path: Path, ref: str) -> Path:
+        if cached_path := self._local_ref_path_cache.get(path):
+            return cached_path
+
         base_path = self.base_path.resolve()
         resolved_path = path.resolve()
         if resolved_path.is_relative_to(base_path) or self.allow_remote_refs is True:
+            self._local_ref_path_cache[path] = resolved_path
             return resolved_path
 
         details = (
@@ -5506,6 +5519,7 @@ def _parse_raw_or_validated_obj(
         self._check_version_specific_features(raw, path)
 
         obj = validated_obj if validated_obj is not None else self._validate_schema_object(raw, path)
+        self._cache_ref_data_type_facts(self.model_resolver.join_path(tuple(path)), obj)
         # Build $recursiveAnchor / $dynamicAnchor indexes for this schema
         self._build_anchor_indexes(obj, path)
         self.parse_obj(name, obj, path)
@@ -5953,6 +5967,7 @@ def _parse_file(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
             with self.root_id_context(raw):
                 # parse $id before parsing $ref
                 root_obj = self._validate_schema_object(raw, path_parts or ["#"])
+                self._cache_ref_data_type_facts(self.model_resolver.join_path(tuple(path_parts or ["#"])), root_obj)
                 self.parse_id(root_obj, [*path_parts, "#"] if path_parts else ["#"])
                 root_key = tuple(path_parts)
                 if root_obj.recursiveAnchor:
@@ -5983,6 +5998,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._cache_ref_data_type_facts(self.model_resolver.join_path(tuple(definition_path)), obj)
                     self.parse_id(obj, definition_path)
                     if obj.recursiveAnchor:
                         ref_path = self._anchor_ref_path(root_key, definition_path)

tests/parser/test_jsonschema.py

@@ -129,6 +129,65 @@ def test_get_x_python_import_path_handles_empty_and_incomplete_metadata() -> Non
     assert parser._get_x_python_import_path({"module": "os", "name": "PathLike"}) == "os.PathLike"
 
 
+def test_get_ref_data_type_uses_cached_validated_definition_facts(mocker: MockerFixture) -> None:
+    """Use facts from the already validated definition instead of validating the ref target again."""
+    parser = JsonSchemaParser(
+        json.dumps({
+            "type": "object",
+            "properties": {"user": {"$ref": "#/$defs/User"}},
+            "required": ["user"],
+            "$defs": {
+                "User": {
+                    "type": "object",
+                    "nullable": True,
+                    "properties": {"name": {"type": "string"}},
+                    "required": ["name"],
+                },
+            },
+        }),
+        strict_nullable=True,
+    )
+    load_ref_schema_object = mocker.spy(parser, "_load_ref_schema_object")
+
+    parser.parse(format_=False)
+
+    load_ref_schema_object.assert_not_called()
+    assert parser._ref_data_type_facts["#/$defs/User"] == (None, True)
+    assert "user: Optional[User]" in dump_templates(list(parser.results))
+
+
+def test_get_ref_data_type_falls_back_when_facts_are_not_cached(mocker: MockerFixture) -> None:
+    """Keep the validation fallback for refs that were not parsed and cached first."""
+    parser = JsonSchemaParser("")
+    parser.raw_obj = {"$defs": {"User": {"type": "object"}}}
+    load_ref_schema_object = mocker.spy(parser, "_load_ref_schema_object")
+
+    parser.get_ref_data_type("#/$defs/User")
+
+    load_ref_schema_object.assert_called_once_with("#/$defs/User")
+
+
+def test_resolve_local_ref_path_caches_safe_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Avoid resolving the same local ref path repeatedly after it has passed safety checks."""
+    parser = JsonSchemaParser(tmp_path / "schema.json")
+    target = tmp_path / "schema.json"
+    original_resolve = Path.resolve
+    calls: list[Path] = []
+
+    def resolve(path: Path, *args: Any, **kwargs: Any) -> Path:
+        calls.append(path)
+        return original_resolve(path, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "resolve", resolve)
+
+    parser._resolve_local_ref_path(target, "schema.json")
+    first_call_count = len(calls)
+    assert first_call_count > 0
+    parser._resolve_local_ref_path(target, "schema.json")
+
+    assert len(calls) == first_call_count
+
+
 def test_json_schema_directory_input_reads_each_source_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     """Test directory input source text is materialized once per parse."""
     user_path, pet_path = _write_simple_json_schemas(tmp_path)