From 15bb0618067d1d34a08b04e8a9e61745c0166096 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Sat, 11 Jul 2026 09:12:17 +0900 Subject: [PATCH 1/2] Retain model import caches --- src/datamodel_code_generator/parser/base.py | 155 +++++++++++-- .../main/builtin_import_cache_retention.py | 13 ++ .../main/builtin_import_cache_retention.txt | 2 + .../custom_import_cache_alias_invalidation.py | 12 + .../__init__.py | 0 .../inner_folder/__init__.py | 0 .../artificial_folder/__init__.py | 0 .../inner_folder/artificial_folder/type_1.py | 7 + .../inner_folder/schema.py | 18 ++ .../inner_folder/type_2.py | 11 + .../subfolder/__init__.py | 0 .../subfolder/type_5.py | 7 + .../type_4.py | 7 + ...stom_import_cache_cross_module_history.txt | 2 + ...custom_import_cache_cross_module_order.txt | 9 + .../main/custom_import_cache_generic_base.py | 43 ++++ .../__init__.py | 0 .../nested_state1.py | 8 + .../nested_state2.py | 8 + .../result1.py | 9 + .../result2.py | 9 + .../state.py | 8 + ...om_import_cache_inherited_enum_history.txt | 6 + ...stom_import_cache_inherited_enum_order.txt | 6 + .../main/custom_import_cache_invalidation.py | 14 ++ tests/main/test_main_general.py | 219 +++++++++++++++++- 26 files changed, 548 insertions(+), 25 deletions(-) create mode 100644 tests/data/expected/main/builtin_import_cache_retention.py create mode 100644 tests/data/expected/main/builtin_import_cache_retention.txt create mode 100644 tests/data/expected/main/custom_import_cache_alias_invalidation.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/__init__.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/inner_folder/__init__.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/__init__.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/type_1.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/inner_folder/schema.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/inner_folder/type_2.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/subfolder/__init__.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/subfolder/type_5.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module/type_4.py create mode 100644 tests/data/expected/main/custom_import_cache_cross_module_history.txt create mode 100644 tests/data/expected/main/custom_import_cache_cross_module_order.txt create mode 100644 tests/data/expected/main/custom_import_cache_generic_base.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/__init__.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/nested_state1.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/nested_state2.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/result1.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/result2.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum/state.py create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum_history.txt create mode 100644 tests/data/expected/main/custom_import_cache_inherited_enum_order.txt create mode 100644 tests/data/expected/main/custom_import_cache_invalidation.py diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index babc261df..0129ef216 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -105,6 +105,7 @@ _PYDANTIC_V2_MODULE: Final = "datamodel_code_generator.model.pydantic_v2" _PYDANTIC_V2_ROOT_MODEL_MODULE: Final = "datamodel_code_generator.model.pydantic_v2.root_model" _TYPED_DICT_MODULE: Final = "datamodel_code_generator.model.typed_dict" +_MODEL_MODULE_PREFIX: Final = "datamodel_code_generator.model." _CLASS_NAME_SEPARATOR_PATTERN: Final = re.compile(r"[^A-Za-z0-9]+") @@ -543,8 +544,9 @@ def _needs_validate_default(data_type: DataType) -> bool: def _alias_base_class_imports( model: DataModel, aliased_imports: dict[tuple[str | None, str], Import], -) -> None: +) -> bool: """Apply aliased imports to a model's base classes and their _additional_imports.""" + changed = False for base_class in model.base_classes: if not base_class.import_: continue @@ -561,6 +563,8 @@ def _alias_base_class_imports( ): # pragma: no branch model._additional_imports[i] = aliased_import # noqa: SLF001 break + changed = True + return changed def _clear_model_imports_cache(models: Iterable[DataModel]) -> None: @@ -571,6 +575,28 @@ def _clear_model_imports_cache(models: Iterable[DataModel]) -> None: model.clear_imports_cache() +def _clear_model_imports_cache_if_retained(model: DataModel, *, can_retain_cache: bool) -> None: + """Clear a changed built-in model without invoking external cache hooks.""" + if can_retain_cache and type(model).__module__.startswith(_MODEL_MODULE_PREFIX): + model.clear_imports_cache() + + +def _can_retain_model_imports_cache( + models: list[DataModel], + *, + configured_types_are_builtin: bool, +) -> bool: + """Return whether all configured and materialized generation types are built in.""" + if not configured_types_are_builtin: + return False + + for model in models: + if type(model).__module__.startswith(_MODEL_MODULE_PREFIX): + continue + return False + return True + + ReferenceMapSet = dict[str, set[str]] SortedDataModels = dict[str, DataModel] @@ -1306,6 +1332,19 @@ def __init__( # noqa: PLR0912, PLR0915 self.data_model_root_type ) self.data_model_field_type: type[DataModelFieldBase] = config.data_model_field_type + self._configured_generation_types_are_builtin = all( + generation_type.__module__.startswith(_MODEL_MODULE_PREFIX) + for generation_type in ( + self.data_model_type, + self.data_model_root_type, + self.data_model_field_type, + type(self.data_type_manager), + ) + ) and all( + (generation_type := getattr(config, attribute, None)) is None + or generation_type.__module__.startswith(_MODEL_MODULE_PREFIX) + for attribute in ("data_model_scalar_type", "data_model_union_type") + ) self.imports: Imports = Imports(config.use_exact_imports) self.use_exact_imports: bool = config.use_exact_imports @@ -2010,6 +2049,8 @@ def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 self, models: list[DataModel], imports: Imports, + *, + can_retain_cache: bool, ) -> None: for model in models: # noqa: PLR1702 for field in model.fields: @@ -2029,6 +2070,7 @@ def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 if has_any_variant: # pragma: no cover field.extras.pop("discriminator", None) field.data_type.discriminator = None + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) continue for data_type in field.data_type.data_types: if not data_type.reference: # pragma: no cover @@ -2103,6 +2145,9 @@ def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 _add_msgspec_base_class_kwarg(discriminator_model, "tag_field", f"'{field_name}'") _add_msgspec_base_class_kwarg(discriminator_model, "tag", repr(expected_value)) discriminator_field.extras["is_classvar"] = True + _clear_model_imports_cache_if_retained( + discriminator_model, can_retain_cache=can_retain_cache + ) # Found the discriminator field, no need to keep looking break @@ -2112,6 +2157,9 @@ def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 _add_msgspec_base_class_kwarg(discriminator_model, "tag_field", f"'{field_name}'") _add_msgspec_base_class_kwarg(discriminator_model, "tag", repr(const_value)) discriminator_field.extras["is_classvar"] = True + _clear_model_imports_cache_if_retained( + discriminator_model, can_retain_cache=can_retain_cache + ) break enum_source = discriminator_field.data_type.find_source(Enum) @@ -2185,7 +2233,7 @@ def _create_set_from_list(cls, data_type: DataType) -> DataType | None: return data_type return None # pragma: no cover - def __replace_unique_list_to_set(self, models: list[DataModel]) -> None: + def __replace_unique_list_to_set(self, models: list[DataModel], *, can_retain_cache: bool) -> None: changed = False for model in models: for model_field in model.fields: @@ -2207,7 +2255,7 @@ def __replace_unique_list_to_set(self, models: list[DataModel]) -> None: model_field.default = converted_default self.generation_store.replace_field_type(model_field, set_data_type) changed = True - if changed: + if changed and not can_retain_cache: _clear_model_imports_cache(models) @classmethod @@ -2237,7 +2285,12 @@ def __mark_set_item_models_hashable(cls, models: list[DataModel]) -> None: class_body_lines.append("__hash__ = object.__hash__") @classmethod - def __set_reference_default_value_to_field(cls, models: list[DataModel]) -> None: + def __set_reference_default_value_to_field( + cls, + models: list[DataModel], + *, + can_retain_cache: bool, + ) -> None: for model in models: for model_field in model.fields: if not model_field.data_type.reference or model_field.has_default: @@ -2248,6 +2301,7 @@ def __set_reference_default_value_to_field(cls, models: list[DataModel]) -> None ): # pragma: no cover model_field.default = model_field.data_type.reference.source.default + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) def __reuse_model(self, models: list[DataModel], require_update_action_models: list[str]) -> None: if not self.reuse_model or self.reuse_scope == ReuseScope.Tree: @@ -2555,10 +2609,12 @@ def __collapse_root_models( # noqa: PLR0912, PLR0914, PLR0915 def __set_default_enum_member( self, models: list[DataModel], + *, + can_retain_cache: bool, ) -> None: if not self.set_default_enum_member: return - for _, model_field, data_type in iter_models_field_data_types(models): + for model, model_field, data_type in iter_models_field_data_types(models): if model_field.default is None: continue if data_type.reference and isinstance(data_type.reference.source, Enum): # pragma: no cover @@ -2571,6 +2627,7 @@ def __set_default_enum_member( if not enum_member: continue model_field.default = enum_member + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) if data_type.alias: if isinstance(enum_member, list): for enum_member_ in enum_member: @@ -2581,6 +2638,8 @@ def __set_default_enum_member( def __set_validate_default_on_fields( # noqa: PLR6301 self, models: list[DataModel], + *, + can_retain_cache: bool, ) -> None: """Set validate_default=True on fields with structured defaults needing validation.""" for model in models: @@ -2593,13 +2652,18 @@ def __set_validate_default_on_fields( # noqa: PLR6301 continue if isinstance(model_field.default, Member): continue + if can_retain_cache and model_field.extras.get("validate_default") is True: + continue if not _needs_validate_default(model_field.data_type): continue model_field.extras["validate_default"] = True + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) def __override_required_field( self, models: list[DataModel], + *, + can_retain_cache: bool = False, ) -> None: changed = False for model in models: @@ -2643,7 +2707,7 @@ def __override_required_field( self.generation_store.insert_field(model, index, copied_original_field) self.generation_store.remove_field(model, model_field) changed = True - if changed: + if changed and not can_retain_cache: _clear_model_imports_cache(models) def __sort_models( @@ -2666,6 +2730,8 @@ def __sort_models( def __change_field_name( self, models: list[DataModel], + *, + can_retain_cache: bool, ) -> None: if not self.data_model_type.SUPPORTS_FIELD_RENAMING: return @@ -2707,6 +2773,7 @@ def __change_field_name( if filed_name != new_filed_name: field.alias = filed_name field.name = new_filed_name + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) if (current_name := field.name) in self.builtin_names and any( _is_builtin_type_collision(current_name, dt) for dt in field.data_type.all_data_types @@ -2714,8 +2781,9 @@ def __change_field_name( if field.alias is None: field.alias = filed_name field.name = f"{current_name}_" + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) - def __set_one_literal_on_default(self, models: list[DataModel]) -> None: + def __set_one_literal_on_default(self, models: list[DataModel], *, can_retain_cache: bool) -> None: if not self.use_one_literal_as_default: return for model in models: @@ -2726,6 +2794,7 @@ def __set_one_literal_on_default(self, models: list[DataModel]) -> None: model_field.required = False if model_field.nullable is not True: # pragma: no cover model_field.nullable = False + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) def __fix_dataclass_field_ordering(self, models: list[DataModel]) -> None: """Fix field ordering for dataclasses with inheritance after defaults are set.""" @@ -2900,6 +2969,7 @@ def __update_type_aliases( pydantic_v2_root_model_type: type[DataModel] | None, *, use_deferred_annotations: bool = True, + can_retain_cache: bool, ) -> None: """Update type aliases and RootModels to properly handle forward references per PEP 484. @@ -2950,6 +3020,7 @@ def __update_type_aliases( if has_forward_ref: model.has_forward_reference = True + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) @classmethod def __postprocess_result_modules(cls, results: dict[tuple[str, ...], Result]) -> dict[tuple[str, ...], Result]: @@ -3007,6 +3078,8 @@ def __alias_shadowed_imports( # noqa: PLR6301 self, models: list[DataModel], all_model_field_names: set[str], + *, + can_retain_cache: bool, ) -> None: aliased_imports: dict[tuple[str | None, str], Import] = {} for _, _model_field, data_type in iter_models_field_data_types(models): @@ -3023,17 +3096,20 @@ def __alias_shadowed_imports( # noqa: PLR6301 if not aliased_imports: return - for _, _model_field, data_type in iter_models_field_data_types(models): + for model, _model_field, data_type in iter_models_field_data_types(models): if data_type and data_type.import_: key = (data_type.import_.from_, data_type.import_.import_) if key in aliased_imports: aliased_import = aliased_imports[key] data_type.type = aliased_import.alias data_type.import_ = aliased_import + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) for model in models: - _alias_base_class_imports(model, aliased_imports) - _clear_model_imports_cache(models) + if _alias_base_class_imports(model, aliased_imports): + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) + if not can_retain_cache: + _clear_model_imports_cache(models) def __apply_generic_base_class( # noqa: PLR0912, PLR0914, PLR0915 self, @@ -3100,13 +3176,26 @@ def __apply_generic_base_class( # noqa: PLR0912, PLR0914, PLR0915 original_base_class = self.data_model_type.BASE_CLASS original_import = Import.from_full_path(original_base_class) if original_base_class else None first_root_module_name = ".".join(first_root_module[:-1]) if first_root_module else "" - for module, _models, target_models, imports in modules_with_targets: + for module, models, target_models, imports in modules_with_targets: current_module_name = ".".join(module[:-1]) if module else "" is_first_root = module == first_root_module + can_retain_cache = _can_retain_model_imports_cache( + models, + configured_types_are_builtin=self._configured_generation_types_are_builtin, + ) for model in target_models: if original_import: # pragma: no branch additional_imports = model._additional_imports # noqa: SLF001 - model._additional_imports = [i for i in additional_imports if i != original_import] # noqa: SLF001 + if can_retain_cache: + if original_import in additional_imports: # pragma: no branch + model._additional_imports = [ # noqa: SLF001 + i for i in additional_imports if i != original_import + ] + model.clear_imports_cache() + else: + model._additional_imports = [ # noqa: SLF001 + i for i in additional_imports if i != original_import + ] parent_refs = [bc.reference for bc in model.base_classes if bc.reference] has_target_model_parent = any(ref.source in all_target_models for ref in parent_refs) if has_target_model_parent: @@ -3697,13 +3786,17 @@ def _process_single_module( # noqa: PLR0913, PLR0917 """Process a single module and return its context.""" imports = Imports(self.use_exact_imports) module, is_init = _resolve_module_file(module_, results) + can_retain_cache = _can_retain_model_imports_cache( + models, + configured_types_are_builtin=self._configured_generation_types_are_builtin, + ) all_module_fields = {field.name for model in models for field in model.fields if field.name is not None} scoped_model_resolver = ModelResolver(exclude_names=all_module_fields) - self.__alias_shadowed_imports(models, all_module_fields) - self.__override_required_field(models) - self.__replace_unique_list_to_set(models) + self.__alias_shadowed_imports(models, all_module_fields, can_retain_cache=can_retain_cache) + self.__override_required_field(models, can_retain_cache=can_retain_cache) + self.__replace_unique_list_to_set(models, can_retain_cache=can_retain_cache) self.__change_from_import( models, imports, @@ -3713,22 +3806,30 @@ def _process_single_module( # noqa: PLR0913, PLR0917 model_path_to_module_name=model_path_to_module_name, ) self.__extract_inherited_enum(models) - self.__set_reference_default_value_to_field(models) + can_retain_cache = can_retain_cache and _can_retain_model_imports_cache( + models, + configured_types_are_builtin=self._configured_generation_types_are_builtin, + ) + self.__set_reference_default_value_to_field(models, can_retain_cache=can_retain_cache) self.__reuse_model(models, require_update_action_models) self.__collapse_root_models(models, unused_models, imports, scoped_model_resolver, model_path_to_module_name) - self.__set_default_enum_member(models) + self.__set_default_enum_member(models, can_retain_cache=can_retain_cache) self.__sort_models(models, imports, use_deferred_annotations=config.use_deferred_annotations) - self.__change_field_name(models) - self.__apply_discriminator_type(models, imports) - self.__set_one_literal_on_default(models) + self.__change_field_name(models, can_retain_cache=can_retain_cache) + self.__apply_discriminator_type(models, imports, can_retain_cache=can_retain_cache) + self.__set_one_literal_on_default(models, can_retain_cache=can_retain_cache) self.__fix_dataclass_field_ordering(models) models = self.__remove_overridden_models(models) self.__apply_type_overrides(models) self.__update_type_aliases( - models, self.pydantic_v2_root_model_type, use_deferred_annotations=config.use_deferred_annotations + models, + self.pydantic_v2_root_model_type, + use_deferred_annotations=config.use_deferred_annotations, + can_retain_cache=can_retain_cache, ) - self.__set_validate_default_on_fields(models) - _clear_model_imports_cache(models) + self.__set_validate_default_on_fields(models, can_retain_cache=can_retain_cache) + if not can_retain_cache: + _clear_model_imports_cache(models) return ModuleContext(module, module_, models, is_init, imports, scoped_model_resolver) @@ -3773,7 +3874,13 @@ def _finalize_modules( for ctx in contexts: self.__change_imported_model_name(ctx.models, ctx.imports, ctx.scoped_model_resolver) - self.__set_validate_default_on_fields(ctx.models) + self.__set_validate_default_on_fields( + ctx.models, + can_retain_cache=_can_retain_model_imports_cache( + ctx.models, + configured_types_are_builtin=self._configured_generation_types_are_builtin, + ), + ) def _generate_module_output( # noqa: PLR0913, PLR0917 self, diff --git a/tests/data/expected/main/builtin_import_cache_retention.py b/tests/data/expected/main/builtin_import_cache_retention.py new file mode 100644 index 000000000..d9d7a67c5 --- /dev/null +++ b/tests/data/expected/main/builtin_import_cache_retention.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel, Field + + +class TestObject(BaseModel): + test_string: Optional[str] = None + + +class Test(BaseModel): + TestObject_1: Optional[TestObject] = Field(None, alias='TestObject', title='TestObject') diff --git a/tests/data/expected/main/builtin_import_cache_retention.txt b/tests/data/expected/main/builtin_import_cache_retention.txt new file mode 100644 index 000000000..05af99d7b --- /dev/null +++ b/tests/data/expected/main/builtin_import_cache_retention.txt @@ -0,0 +1,2 @@ +TestObject:retained +Test:invalidated diff --git a/tests/data/expected/main/custom_import_cache_alias_invalidation.py b/tests/data/expected/main/custom_import_cache_alias_invalidation.py new file mode 100644 index 000000000..787bfc581 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_alias_invalidation.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from datetime import date as date_aliased +from typing import Optional + +from pydantic import AwareDatetime, BaseModel + + +class TypeDate(BaseModel): + cache_clear_history = ['empty', 'cached', 'cached'] + date: Optional[date_aliased] + datetime: Optional[AwareDatetime] = None diff --git a/tests/data/expected/main/custom_import_cache_cross_module/__init__.py b/tests/data/expected/main/custom_import_cache_cross_module/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/__init__.py b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/__init__.py b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/type_1.py b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/type_1.py new file mode 100644 index 000000000..cefdec216 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/artificial_folder/type_1.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from msgspec import Struct + + +class Type1(Struct, tag_field='type_', tag='a'): + pass diff --git a/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/schema.py b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/schema.py new file mode 100644 index 000000000..3d592c5fc --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/schema.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import Union + +from msgspec import Struct + +from .. import type_4 +from ..subfolder import type_5 +from . import type_2 +from .artificial_folder import type_1 + + +class Type3(Struct, tag_field='type_', tag='c'): + pass + + +class Response(Struct): + inner: Union[type_1.Type1, type_2.Type2, Type3, type_4.Type4, type_5.Type5] diff --git a/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/type_2.py b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/type_2.py new file mode 100644 index 000000000..33abefe1c --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module/inner_folder/type_2.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +from typing import Union + +from msgspec import UNSET, Struct, UnsetType + +from .artificial_folder import type_1 + + +class Type2(Struct, tag_field='type_', tag='b'): + ref_type: Union[type_1.Type1, UnsetType] = UNSET diff --git a/tests/data/expected/main/custom_import_cache_cross_module/subfolder/__init__.py b/tests/data/expected/main/custom_import_cache_cross_module/subfolder/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/expected/main/custom_import_cache_cross_module/subfolder/type_5.py b/tests/data/expected/main/custom_import_cache_cross_module/subfolder/type_5.py new file mode 100644 index 000000000..4214ce6f0 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module/subfolder/type_5.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from msgspec import Struct + + +class Type5(Struct, tag_field='type_', tag='e'): + pass diff --git a/tests/data/expected/main/custom_import_cache_cross_module/type_4.py b/tests/data/expected/main/custom_import_cache_cross_module/type_4.py new file mode 100644 index 000000000..82584955c --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module/type_4.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from msgspec import Struct + + +class Type4(Struct, tag_field='type_', tag='d'): + pass diff --git a/tests/data/expected/main/custom_import_cache_cross_module_history.txt b/tests/data/expected/main/custom_import_cache_cross_module_history.txt new file mode 100644 index 000000000..837def637 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module_history.txt @@ -0,0 +1,2 @@ +?:empty +Type1:cached diff --git a/tests/data/expected/main/custom_import_cache_cross_module_order.txt b/tests/data/expected/main/custom_import_cache_cross_module_order.txt new file mode 100644 index 000000000..89ba0d97c --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_cross_module_order.txt @@ -0,0 +1,9 @@ +inner_folder/__init__.py +inner_folder/artificial_folder/__init__.py +subfolder/__init__.py +__init__.py +inner_folder/artificial_folder/type_1.py +subfolder/type_5.py +inner_folder/type_2.py +inner_folder/schema.py +type_4.py diff --git a/tests/data/expected/main/custom_import_cache_generic_base.py b/tests/data/expected/main/custom_import_cache_generic_base.py new file mode 100644 index 000000000..e05ee8d38 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_generic_base.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel as _BaseModel +from pydantic import ConfigDict + + +class BaseModel(_BaseModel): + model_config = ConfigDict( + extra='forbid', + ) + + +class Foo(BaseModel): + model_config = ConfigDict( + extra='allow', + ) + cache_clear_history = ['empty', 'cached', 'empty'] + x: Optional[int] = None + + +class Bar(BaseModel): + model_config = ConfigDict( + extra='forbid', + ) + cache_clear_history = ['empty', 'cached', 'empty'] + y: Optional[int] = None + + +class Baz(BaseModel): + cache_clear_history = ['empty', 'cached', 'empty'] + z: Optional[int] = None + + +class Test(BaseModel): + model_config = ConfigDict( + extra='forbid', + ) + cache_clear_history = ['empty', 'cached', 'empty'] + foo: Foo + bar: Optional[Bar] = None + baz: Optional[Baz] = None diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/__init__.py b/tests/data/expected/main/custom_import_cache_inherited_enum/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state1.py b/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state1.py new file mode 100644 index 000000000..78b8d9a67 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state1.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from enum import Enum + + +class NestedState1(Enum): + field_1 = '1' + field_2 = '2' diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state2.py b/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state2.py new file mode 100644 index 000000000..3c0c4c229 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum/nested_state2.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from enum import Enum + + +class NestedState2(Enum): + field_1 = '1' + field_2 = '2' diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/result1.py b/tests/data/expected/main/custom_import_cache_inherited_enum/result1.py new file mode 100644 index 000000000..193281267 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum/result1.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +from pydantic import BaseModel + +from . import nested_state1 + + +class Result1(BaseModel): + state: nested_state1.NestedState1 diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/result2.py b/tests/data/expected/main/custom_import_cache_inherited_enum/result2.py new file mode 100644 index 000000000..320f23e84 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum/result2.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +from pydantic import BaseModel + +from . import nested_state2 + + +class Result2(BaseModel): + state: nested_state2.NestedState2 diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum/state.py b/tests/data/expected/main/custom_import_cache_inherited_enum/state.py new file mode 100644 index 000000000..85c347d5a --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum/state.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from enum import Enum + + +class State(Enum): + field_1 = '1' + field_2 = '2' diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum_history.txt b/tests/data/expected/main/custom_import_cache_inherited_enum_history.txt new file mode 100644 index 000000000..8d4f2fbd9 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum_history.txt @@ -0,0 +1,6 @@ +?:empty +State:cached +?:empty +NestedState2:empty +?:empty +NestedState1:empty diff --git a/tests/data/expected/main/custom_import_cache_inherited_enum_order.txt b/tests/data/expected/main/custom_import_cache_inherited_enum_order.txt new file mode 100644 index 000000000..237c274d4 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_inherited_enum_order.txt @@ -0,0 +1,6 @@ +__init__.py +state.py +result2.py +result1.py +nested_state2.py +nested_state1.py diff --git a/tests/data/expected/main/custom_import_cache_invalidation.py b/tests/data/expected/main/custom_import_cache_invalidation.py new file mode 100644 index 000000000..16b643df6 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_invalidation.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from typing import Any, List, Optional + +from pydantic import BaseModel, Field, conint + + +class Person(BaseModel): + cache_clear_history = ['empty', 'cached'] + firstName: Optional[str] = Field(None, description="The person's first name.") + lastName: Optional[str] = Field(None, description="The person's last name.") + age: Optional[conint(ge=0)] = Field(None, description='Age in years which must be equal to or greater than zero.') + friends: Optional[List[Any]] = None + comment: None = Field(None) diff --git a/tests/main/test_main_general.py b/tests/main/test_main_general.py index 7c8c449f7..036ceffab 100644 --- a/tests/main/test_main_general.py +++ b/tests/main/test_main_general.py @@ -11,7 +11,7 @@ from argparse import ArgumentTypeError, BooleanOptionalAction, Namespace from collections import defaultdict from io import StringIO -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import black import pytest @@ -55,6 +55,7 @@ assert_generated_modules_output, assert_httpx_get_kwargs, assert_no_uncommented_generated_code, + assert_output, assert_runtime_import_package, assert_warnings_contain, assert_warnings_do_not_contain, @@ -93,6 +94,222 @@ class _GenerateParseAbort(BaseException): """Test-only parse abort that is not an Exception subclass.""" +def test_parser_retains_builtin_import_cache_and_invalidates_custom_cache() -> None: + """Retain built-in caches while keeping the legacy custom-model invalidation contract.""" + from datamodel_code_generator.model import DataModel, get_data_model_types + from datamodel_code_generator.model.pydantic_v2 import BaseModel + from datamodel_code_generator.parser.jsonschema import JsonSchemaParser + + class CacheProbeJsonSchemaParser(JsonSchemaParser): + cache_reuse_manifest: tuple[str, ...] = () + + def _process_single_module(self, module_: Any, models: list[Any], *args: Any, **kwargs: Any) -> Any: + for model in models: + _ = model.imports + cached_imports = tuple(model.__dict__[model._IMPORTS_CACHE_KEY] for model in models) + context = super()._process_single_module(module_, models, *args, **kwargs) + self.cache_reuse_manifest += tuple( + f"{model.class_name}:" + f"{'retained' if cached is model.__dict__.get(model._IMPORTS_CACHE_KEY) else 'invalidated'}" + for cached, model in zip(cached_imports, context.models, strict=True) + ) + return context + + class CacheAwareBaseModel(BaseModel): + def clear_imports_cache(self) -> None: + cache_state = "cached" if self._IMPORTS_CACHE_KEY in self.__dict__ else "empty" + history = self.__dict__.setdefault("cache_clear_history", []) + history.append(cache_state) + if (extra_template_data := getattr(self, "extra_template_data", None)) is not None: + extra_template_data["class_body_lines"] = [f"cache_clear_history = {history!r}"] + super().clear_imports_cache() + + class InjectingJsonSchemaParser(JsonSchemaParser): + def _create_data_model( + self, + model_type: type[DataModel] | None = None, + **kwargs: Any, + ) -> DataModel: + if model_type is None or model_type is self.data_model_type: + model_type = CacheAwareBaseModel + return super()._create_data_model(model_type, **kwargs) + + model_types = get_data_model_types( + DataModelType.PydanticV2BaseModel, + target_python_version=PythonVersion.PY_311, + ) + input_path = JSON_SCHEMA_DATA_PATH / "field_has_same_name.json" + parser_options = { + "base_path": input_path.parent, + "data_model_root_type": model_types.root_model, + "data_model_field_type": model_types.field_model, + "data_type_manager_type": model_types.data_type_manager, + "dump_resolve_reference_action": model_types.dump_resolve_reference_action, + "formatters": [Formatter.BUILTIN], + "target_python_version": PythonVersion.PY_311, + } + parser = CacheProbeJsonSchemaParser(input_path, **parser_options) + assert_output(parser.parse(), EXPECTED_MAIN_PATH / "builtin_import_cache_retention.py") + assert_output( + "\n".join(parser.cache_reuse_manifest) + "\n", + EXPECTED_MAIN_PATH / "builtin_import_cache_retention.txt", + ) + + input_path = JSON_SCHEMA_DATA_PATH / "person.json" + custom_parser = JsonSchemaParser( + input_path, + data_model_type=CacheAwareBaseModel, + **parser_options, + ) + assert_output( + custom_parser.parse(), + EXPECTED_MAIN_PATH / "custom_import_cache_invalidation.py", + ) + + injected_parser = InjectingJsonSchemaParser(input_path, **parser_options) + assert_output( + injected_parser.parse(), + EXPECTED_MAIN_PATH / "custom_import_cache_invalidation.py", + ) + + alias_input_path = JSON_SCHEMA_DATA_PATH / "alias_import_alias" / "date.schema.json" + alias_parser = JsonSchemaParser( + alias_input_path, + data_model_type=CacheAwareBaseModel, + **{**parser_options, "base_path": alias_input_path.parent}, + ) + assert_output( + alias_parser.parse(), + EXPECTED_MAIN_PATH / "custom_import_cache_alias_invalidation.py", + ) + + generic_input_path = JSON_SCHEMA_DATA_PATH / "extra_fields.json" + generic_parser = JsonSchemaParser( + generic_input_path, + data_model_type=CacheAwareBaseModel, + extra_fields="forbid", + use_generic_base_class=True, + **{**parser_options, "base_path": generic_input_path.parent}, + ) + assert_output( + generic_parser.parse(), + EXPECTED_MAIN_PATH / "custom_import_cache_generic_base.py", + ) + + +def test_parser_preserves_cross_module_external_import_cache_hook() -> None: + """Do not invoke an external model's cache hook from another built-in module.""" + from datamodel_code_generator.model import DataModel, get_data_model_types + from datamodel_code_generator.model.msgspec import Struct + from datamodel_code_generator.parser.jsonschema import JsonSchemaParser + + class ProbeStruct(Struct): + clear_history: ClassVar[list[str]] = [] + + def clear_imports_cache(self) -> None: + state = "cached" if self._IMPORTS_CACHE_KEY in self.__dict__ else "empty" + name = getattr(self, "class_name", "") or "?" + self.clear_history.append(f"{name}:{state}") + super().clear_imports_cache() + + class SelectiveExternalStructParser(JsonSchemaParser): + def _create_data_model( + self, + model_type: type[DataModel] | None = None, + **kwargs: Any, + ) -> DataModel: + reference = kwargs.get("reference") + if ( + reference is not None + and reference.name == "Type1" + and (model_type is None or model_type is self.data_model_type) + ): + model_type = ProbeStruct + return super()._create_data_model(model_type, **kwargs) + + model_types = get_data_model_types( + DataModelType.MsgspecStruct, + target_python_version=PythonVersion.PY_311, + ) + input_path = JSON_SCHEMA_DATA_PATH / "discriminator_with_external_reference" + parser = SelectiveExternalStructParser( + input_path, + base_path=input_path, + data_model_type=model_types.data_model, + data_model_root_type=model_types.root_model, + data_model_field_type=model_types.field_model, + data_type_manager_type=model_types.data_type_manager, + dump_resolve_reference_action=model_types.dump_resolve_reference_action, + formatters=[Formatter.BUILTIN], + target_python_version=PythonVersion.PY_311, + ) + modules = cast("dict[tuple[str, ...], Any]", parser.parse()) + + assert_generated_modules_output(modules, EXPECTED_MAIN_PATH / "custom_import_cache_cross_module") + assert_output( + "\n".join(ProbeStruct.clear_history) + "\n", + EXPECTED_MAIN_PATH / "custom_import_cache_cross_module_history.txt", + ) + assert_output( + "\n".join("/".join(module) for module in modules) + "\n", + EXPECTED_MAIN_PATH / "custom_import_cache_cross_module_order.txt", + ) + + +def test_parser_rechecks_external_enum_after_module_materialization() -> None: + """Keep the external hook fallback when an inherited enum replaces a built-in wrapper.""" + from datamodel_code_generator import ModuleSplitMode + from datamodel_code_generator.model import get_data_model_types + from datamodel_code_generator.model.enum import Enum + from datamodel_code_generator.parser.openapi import OpenAPIParser, OpenAPIScope + + class ProbeEnum(Enum): + clear_history: ClassVar[list[str]] = [] + + def clear_imports_cache(self) -> None: + state = "cached" if self._IMPORTS_CACHE_KEY in self.__dict__ else "empty" + name = getattr(self, "class_name", "") or "?" + self.clear_history.append(f"{name}:{state}") + super().clear_imports_cache() + + class SelectiveExternalEnumParser(OpenAPIParser): + def _get_enum_model_class(self, type_: Any, enum_values: list[Any]) -> tuple[type[Enum], Any]: + _, remaining_type = super()._get_enum_model_class(type_, enum_values) + return ProbeEnum, remaining_type + + model_types = get_data_model_types( + DataModelType.PydanticV2BaseModel, + target_python_version=PythonVersion.PY_311, + ) + input_path = OPEN_API_DATA_PATH / "nested_enum.json" + parser = SelectiveExternalEnumParser( + input_path, + base_path=input_path.parent, + data_model_type=model_types.data_model, + data_model_root_type=model_types.root_model, + data_model_field_type=model_types.field_model, + data_type_manager_type=model_types.data_type_manager, + dump_resolve_reference_action=model_types.dump_resolve_reference_action, + formatters=[Formatter.BUILTIN], + openapi_scopes=[OpenAPIScope.Schemas], + target_python_version=PythonVersion.PY_311, + ) + modules = cast( + "dict[tuple[str, ...], Any]", + parser.parse(module_split_mode=ModuleSplitMode.Single), + ) + + assert_generated_modules_output(modules, EXPECTED_MAIN_PATH / "custom_import_cache_inherited_enum") + assert_output( + "\n".join(ProbeEnum.clear_history) + "\n", + EXPECTED_MAIN_PATH / "custom_import_cache_inherited_enum_history.txt", + ) + assert_output( + "\n".join("/".join(module) for module in modules) + "\n", + EXPECTED_MAIN_PATH / "custom_import_cache_inherited_enum_order.txt", + ) + + CLI_E2E_COVERED_GENERATE_KWARGS = { "emit_model_metadata", "infer_union_variant_names", From 53e743458e939171eeaa766591949c0e08a63788 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Sat, 11 Jul 2026 13:39:57 +0900 Subject: [PATCH 2/2] Cover import cache mutation edge cases --- src/datamodel_code_generator/parser/base.py | 5 ++- ...in_import_cache_unique_items_unhashable.py | 18 +++++++++ ...n_import_cache_unique_items_unhashable.txt | 3 ++ .../main/custom_import_cache_unique_items.py | 24 ++++++++++++ .../unique_items_unhashable_default.json | 39 +++++++++++++++++++ tests/main/test_main_general.py | 36 ++++++++++++++++- 6 files changed, 121 insertions(+), 4 deletions(-) create mode 100644 tests/data/expected/main/builtin_import_cache_unique_items_unhashable.py create mode 100644 tests/data/expected/main/builtin_import_cache_unique_items_unhashable.txt create mode 100644 tests/data/expected/main/custom_import_cache_unique_items.py create mode 100644 tests/data/jsonschema/unique_items_unhashable_default.json diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index 0129ef216..d6e81a895 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -2249,8 +2249,9 @@ def __replace_unique_list_to_set(self, models: list[DataModel], *, can_retain_ca try: converted_default = set(model_field.default) except TypeError: - # Elements are not hashable (e.g., contains dicts) - # Skip both type and default conversion to keep consistency + # Keep the unhashable default unchanged. Nested list types may + # already have been converted in place, so invalidate their imports. + _clear_model_imports_cache_if_retained(model, can_retain_cache=can_retain_cache) continue model_field.default = converted_default self.generation_store.replace_field_type(model_field, set_data_type) diff --git a/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.py b/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.py new file mode 100644 index 000000000..53430fd17 --- /dev/null +++ b/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional, Set + +from pydantic import BaseModel + + +class Item(BaseModel): + values: Optional[Set[Dict[str, Any]]] = [{'a': 1}] + + +class Left(BaseModel): + item: Optional[Item] = None + + +class Model(BaseModel): + left: Optional[Left] = None + right: Optional[Left] = None diff --git a/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.txt b/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.txt new file mode 100644 index 000000000..7ddb382a7 --- /dev/null +++ b/tests/data/expected/main/builtin_import_cache_unique_items_unhashable.txt @@ -0,0 +1,3 @@ +Item:invalidated +Left:retained +Model:retained diff --git a/tests/data/expected/main/custom_import_cache_unique_items.py b/tests/data/expected/main/custom_import_cache_unique_items.py new file mode 100644 index 000000000..991136bd7 --- /dev/null +++ b/tests/data/expected/main/custom_import_cache_unique_items.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from enum import StrEnum +from typing import Optional, Set + +from pydantic import BaseModel + + +class Status(StrEnum): + active = 'active' + inactive = 'inactive' + pending = 'pending' + + +class Item(BaseModel): + cache_clear_history = ['empty', 'cached', 'cached'] + __hash__ = object.__hash__ + name: Optional[str] = None + + +class Container(BaseModel): + cache_clear_history = ['empty', 'cached', 'empty', 'empty', 'empty', 'empty', 'cached'] + statuses: Optional[Set[Status]] = None + items: Optional[Set[Item]] = None diff --git a/tests/data/jsonschema/unique_items_unhashable_default.json b/tests/data/jsonschema/unique_items_unhashable_default.json new file mode 100644 index 000000000..853dfa902 --- /dev/null +++ b/tests/data/jsonschema/unique_items_unhashable_default.json @@ -0,0 +1,39 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "Model", + "type": "object", + "properties": { + "left": { + "type": "object", + "properties": { + "item": { + "type": "object", + "properties": { + "values": { + "type": ["array", "null"], + "items": {"type": "object"}, + "uniqueItems": true, + "default": [{"a": 1}] + } + } + } + } + }, + "right": { + "type": "object", + "properties": { + "item": { + "type": "object", + "properties": { + "values": { + "type": ["array", "null"], + "items": {"type": "object"}, + "uniqueItems": true, + "default": [{"a": 1}] + } + } + } + } + } + } +} diff --git a/tests/main/test_main_general.py b/tests/main/test_main_general.py index 036ceffab..43826cc60 100644 --- a/tests/main/test_main_general.py +++ b/tests/main/test_main_general.py @@ -130,8 +130,10 @@ def _create_data_model( model_type: type[DataModel] | None = None, **kwargs: Any, ) -> DataModel: - if model_type is None or model_type is self.data_model_type: - model_type = CacheAwareBaseModel + model_type = { + None: CacheAwareBaseModel, + self.data_model_type: CacheAwareBaseModel, + }.get(model_type, model_type) return super()._create_data_model(model_type, **kwargs) model_types = get_data_model_types( @@ -155,6 +157,25 @@ def _create_data_model( EXPECTED_MAIN_PATH / "builtin_import_cache_retention.txt", ) + input_path = JSON_SCHEMA_DATA_PATH / "unique_items_unhashable_default.json" + unhashable_default_parser = CacheProbeJsonSchemaParser( + input_path, + **{ + **parser_options, + "collapse_reuse_models": True, + "reuse_model": True, + "use_unique_items_as_set": True, + }, + ) + assert_output( + unhashable_default_parser.parse(), + EXPECTED_MAIN_PATH / "builtin_import_cache_unique_items_unhashable.py", + ) + assert_output( + "\n".join(unhashable_default_parser.cache_reuse_manifest) + "\n", + EXPECTED_MAIN_PATH / "builtin_import_cache_unique_items_unhashable.txt", + ) + input_path = JSON_SCHEMA_DATA_PATH / "person.json" custom_parser = JsonSchemaParser( input_path, @@ -172,6 +193,17 @@ def _create_data_model( EXPECTED_MAIN_PATH / "custom_import_cache_invalidation.py", ) + input_path = JSON_SCHEMA_DATA_PATH / "unique_items_enum_set.json" + unique_items_parser = JsonSchemaParser( + input_path, + data_model_type=CacheAwareBaseModel, + **{**parser_options, "use_unique_items_as_set": True}, + ) + assert_output( + unique_items_parser.parse(), + EXPECTED_MAIN_PATH / "custom_import_cache_unique_items.py", + ) + alias_input_path = JSON_SCHEMA_DATA_PATH / "alias_import_alias" / "date.schema.json" alias_parser = JsonSchemaParser( alias_input_path,