From 8e9dd00cb319793cd9e65cdfe9338f0d5ee3f37d Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Mon, 22 Jun 2026 02:25:51 +0900 Subject: [PATCH] Reuse validated schema definitions --- .../parser/jsonschema.py | 12 ++- .../jsonschema/same_file_ref_fast_path.py | 15 ++++ .../jsonschema/same_file_ref_fast_path.json | 25 ++++++ tests/main/jsonschema/test_main_jsonschema.py | 10 +++ tests/parser/test_jsonschema.py | 80 ++++++++++++++++++- 5 files changed, 140 insertions(+), 2 deletions(-) create mode 100644 tests/data/expected/main/jsonschema/same_file_ref_fast_path.py create mode 100644 tests/data/jsonschema/same_file_ref_fast_path.json diff --git a/src/datamodel_code_generator/parser/jsonschema.py b/src/datamodel_code_generator/parser/jsonschema.py index 45ab23310..4cb8c6833 100644 --- a/src/datamodel_code_generator/parser/jsonschema.py +++ b/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) diff --git a/tests/data/expected/main/jsonschema/same_file_ref_fast_path.py b/tests/data/expected/main/jsonschema/same_file_ref_fast_path.py new file mode 100644 index 000000000..e5825571f --- /dev/null +++ b/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 diff --git a/tests/data/jsonschema/same_file_ref_fast_path.json b/tests/data/jsonschema/same_file_ref_fast_path.json new file mode 100644 index 000000000..a615346ce --- /dev/null +++ b/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" + ] + } + } +} diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py index e7d1eb412..1b717a20f 100644 --- a/tests/main/jsonschema/test_main_jsonschema.py +++ b/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( diff --git a/tests/parser/test_jsonschema.py b/tests/parser/test_jsonschema.py index 577d49b66..dbdb7a171 100644 --- a/tests/parser/test_jsonschema.py +++ b/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(