Reuse validated JSON Schema objects (#3466)

* Reuse validated JSON Schema objects

* Cache JSON Schema refs on first load

Author: koxudaxi

src/datamodel_code_generator/model/base.py

@@ -431,13 +431,15 @@ def _can_collect_imports_without_type_hint(self, *, needs_annotated: bool) -> bo
             return False
 
         data_type = self.data_type
+        if data_type.use_serialize_as_any:
+            return False
+
         type_name = data_type.alias or data_type.type
         if type_name and (UNION_PREFIX in type_name or OPTIONAL_PREFIX in type_name):
             return False
 
         return not (
-            data_type.reference
-            or data_type.data_types
+            data_type.data_types
             or data_type.dict_key
             or data_type.literals
             or data_type.enum_member_literals
@@ -472,23 +474,33 @@ def _needs_optional_import_without_type_hint(self) -> bool:
         ):
             return False
 
+        if self._reference_source_nullable_without_type_hint():
+            return True
+
         match self.nullable:
             case True:
                 return True
             case False:
                 return False
-            case None if self.required:
-                return bool(self.type_has_null)
             case None:
-                return bool(self.fall_back_to_nullable)
+                return bool(self.type_has_null) if self.required else bool(self.fall_back_to_nullable)
         return False  # pragma: no cover
 
+    def _reference_source_nullable_without_type_hint(self) -> bool:
+        """Return whether a referenced non-alias model should make this type optional."""
+        reference = self.data_type.reference
+        if reference is None:
+            return False
+        source = reference.source
+        return not getattr(source, "is_alias", False) and bool(getattr(source, "nullable", False))
+
     def _has_renderable_data_type_without_type_hint(self) -> bool:
         """Return whether this simple DataType renders to something other than None."""
         data_type = self.data_type
         return bool(
             data_type.alias
             or data_type.type
+            or data_type.reference
             or data_type.is_dict
             or data_type.is_list
             or data_type.is_set

src/datamodel_code_generator/parser/jsonschema.py

@@ -901,6 +901,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._ref_schema_object_cache: dict[str, JsonSchemaObject] = {}
         self._force_base_model_refs: set[str] = set()
         self._force_base_model_generation = False
         self.field_keys: set[str] = {
@@ -2462,6 +2463,9 @@ def _deep_merge(self, dict1: dict[Any, Any], dict2: dict[Any, Any]) -> dict[Any,
     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)
+        if (cached := self._ref_schema_object_cache.get(resolved_ref)) is not None:
+            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
 
@@ -2470,7 +2474,9 @@ def _load_ref_schema_object(self, ref: str) -> JsonSchemaObject:
             pointer = split_json_pointer(raw_doc, fragment)
             target_schema = get_model_by_path(raw_doc, pointer)
 
-        return self._validate_schema_object(target_schema, [resolved_ref])
+        ref_schema = self._validate_schema_object(target_schema, [resolved_ref])
+        self._ref_schema_object_cache[resolved_ref] = ref_schema
+        return ref_schema
 
     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."""
@@ -5486,14 +5492,24 @@ def parse_raw_obj(
         path: list[str],
     ) -> None:
         """Parse a raw dictionary into a JsonSchemaObject and process it."""
+        self._parse_raw_or_validated_obj(name, raw, path)
+
+    def _parse_raw_or_validated_obj(
+        self,
+        name: str,
+        raw: dict[str, YamlValue] | YamlValue,
+        path: list[str],
+        validated_obj: JsonSchemaObject | None = None,
+    ) -> None:
+        """Parse a raw schema, reusing a validated object when available."""
         if isinstance(raw, dict) and "x-python-import" in raw:
             self._handle_python_import(name, path)
             return
 
         # Strict mode: check for version-specific features before validation
         self._check_version_specific_features(raw, path)
 
-        obj = self._validate_schema_object(raw, path)
+        obj = validated_obj if validated_obj is not None else self._validate_schema_object(raw, path)
         # Build $recursiveAnchor / $dynamicAnchor indexes for this schema
         self._build_anchor_indexes(obj, path)
         self.parse_obj(name, obj, path)
@@ -5964,11 +5980,13 @@ def _parse_file(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
                     *definition_entries,
                 ]
                 seen_definition_metadata_paths: set[tuple[str, ...]] = set()
+                validated_definition_objects: dict[tuple[str, ...], JsonSchemaObject] = {}
                 for _key, model, definition_path in definition_metadata_entries:
                     if (definition_path_key := tuple(definition_path)) in seen_definition_metadata_paths:
                         continue
                     seen_definition_metadata_paths.add(definition_path_key)
                     obj = self._validate_schema_object(model, definition_path)
+                    validated_definition_objects[definition_path_key] = obj
                     self.parse_id(obj, definition_path)
                     if obj.recursiveAnchor:
                         ref_path = self._anchor_ref_path(root_key, definition_path)
@@ -5986,7 +6004,12 @@ def _parse_file(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
                 for key, model, path in definition_entries:
                     reference = self.model_resolver.get(path)
                     if not reference or not reference.loaded:
-                        self.parse_raw_obj(key, model, path)
+                        self._parse_raw_or_validated_obj(
+                            key,
+                            model,
+                            path,
+                            validated_definition_objects.get(tuple(path)),
+                        )
 
                 key = tuple(path_parts)
                 reserved_refs = set(self.reserved_refs.get(key) or [])

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

@@ -0,0 +1,17 @@
+# generated by datamodel-codegen:
+#   filename:  ref_nullable_import_fast_path.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class User(BaseModel):
+    name: str
+
+
+class RefNullableImportFastPath(BaseModel):
+    maybe_user: Optional[User]

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

@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  serialize_as_any_import_fast_path.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel, SerializeAsAny
+
+
+class User(BaseModel):
+    name: str
+
+
+class Admin(User):
+    role: str
+
+
+class SerializeAsAnyImportFastPath(BaseModel):
+    user: SerializeAsAny[User]

tests/data/jsonschema/ref_nullable_import_fast_path.json

@@ -0,0 +1,28 @@
+{
+  "title": "RefNullableImportFastPath",
+  "type": "object",
+  "properties": {
+    "maybe_user": {
+      "$ref": "#/definitions/User"
+    }
+  },
+  "required": [
+    "maybe_user"
+  ],
+  "definitions": {
+    "User": {
+      "type": [
+        "object",
+        "null"
+      ],
+      "properties": {
+        "name": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "name"
+      ]
+    }
+  }
+}

tests/data/jsonschema/serialize_as_any_import_fast_path.json

@@ -0,0 +1,43 @@
+{
+  "title": "SerializeAsAnyImportFastPath",
+  "type": "object",
+  "properties": {
+    "user": {
+      "$ref": "#/definitions/User"
+    }
+  },
+  "required": [
+    "user"
+  ],
+  "definitions": {
+    "User": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "name"
+      ]
+    },
+    "Admin": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/User"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "role": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "role"
+          ]
+        }
+      ]
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2219,6 +2219,33 @@ def test_main_import_fast_path_primitives(output_file: Path) -> None:
     )
 
 
+def test_main_ref_nullable_import_fast_path(output_file: Path) -> None:
+    """Test referenced nullable fields keep Optional imports with no union operator."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "ref_nullable_import_fast_path.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        extra_args=["--no-use-union-operator"],
+    )
+
+
+def test_main_serialize_as_any_import_fast_path(output_file: Path) -> None:
+    """Test SerializeAsAny reference fields keep generated imports."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "serialize_as_any_import_fast_path.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--use-serialize-as-any",
+        ],
+        force_exec_validation=True,
+    )
+
+
 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/model/test_base.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import Any
@@ -59,6 +60,14 @@ class C(DataModel):
     """Test helper class for DataModel testing without template path."""
 
 
+@dataclass
+class ReferenceSource:
+    """Test helper for reference source nullability."""
+
+    nullable: bool
+    is_alias: bool = False
+
+
 template: str = """{%- for decorator in decorators -%}
 {{ decorator }}
 {%- endfor %}
@@ -451,7 +460,6 @@ def test_data_field() -> None:
 @pytest.mark.parametrize(
     "data_type",
     [
-        DataType(reference=Reference(path="Ref", original_name="Ref", name="Ref")),
         DataType(data_types=[DataType(type="str")]),
         DataType(is_dict=True, dict_key=DataType(type="str")),
         DataType(literals=["value"]),
@@ -559,6 +567,24 @@ def test_field_import_fast_path_collects_simple_data_type_imports(
     assert field.imports == expected_imports
 
 
+def test_field_import_fast_path_collects_nullable_reference_import() -> None:
+    """Test reference source nullability contributes Optional without rendering type hints."""
+    reference = Reference(path="#/definitions/User", name="User")
+    reference.source = ReferenceSource(nullable=True)
+    field = DataModelFieldBase(name="user", data_type=DataType(reference=reference), required=True)
+
+    assert field.imports == (IMPORT_OPTIONAL,)
+
+
+def test_field_import_fast_path_ignores_nullable_alias_reference() -> None:
+    """Test nullable aliases do not make the referencing field import Optional."""
+    reference = Reference(path="#/definitions/UserAlias", name="UserAlias")
+    reference.source = ReferenceSource(nullable=True, is_alias=True)
+    field = DataModelFieldBase(name="user", data_type=DataType(reference=reference), required=True)
+
+    assert field.imports == ()
+
+
 def test_field_import_fallback_collects_annotated_import() -> None:
     """Test the fallback path includes Annotated when requested."""
     field = DataModelFieldBase(name="a", data_type=DataType(type="str"), required=True)