From 892b1691f8142ae7077ca9621b778710c8d6f1ed Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Wed, 8 Jul 2026 01:29:03 +0900 Subject: [PATCH] Fix type overrides for base refs --- docs/cli-reference/typing-customization.md | 3 +- docs/llms-full.txt | 3 +- src/datamodel_code_generator/parser/base.py | 57 +++++++++++---- .../main/type_overrides_base_class.py | 32 +++++++++ .../expected/main/type_overrides_dict_key.py | 16 +++++ .../jsonschema/type_overrides_base_class.json | 71 +++++++++++++++++++ .../jsonschema/type_overrides_dict_key.json | 27 +++++++ tests/main/test_main_general.py | 65 ++++++++++++++++- 8 files changed, 257 insertions(+), 17 deletions(-) create mode 100644 tests/data/expected/main/type_overrides_base_class.py create mode 100644 tests/data/expected/main/type_overrides_dict_key.py create mode 100644 tests/data/jsonschema/type_overrides_base_class.json create mode 100644 tests/data/jsonschema/type_overrides_dict_key.json diff --git a/docs/cli-reference/typing-customization.md b/docs/cli-reference/typing-customization.md index df1f22b9f..9b3b225f5 100644 --- a/docs/cli-reference/typing-customization.md +++ b/docs/cli-reference/typing-customization.md @@ -2813,7 +2813,7 @@ instead of generating them. | Format | Description | |--------|-------------| -| `{"ModelName": "package.Type"}` | Model-level: Skip generating `ModelName` and import from `package` | +| `{"ModelName": "package.Type"}` | Model-level: Skip generation; replace field and inheritance refs | | `{"Model.field": "package.Type"}` | Scoped: Override only specific field in specific model | **Common Use Cases:** @@ -2836,6 +2836,7 @@ instead of generating them. !!! note "Model-level overrides skip generation" When you specify a model-level override (without a dot in the key), the generator will **skip generating that model entirely** and import it from the specified package instead. + References to that model are replaced in field annotations and `allOf` inheritance base classes. ??? example "Examples" diff --git a/docs/llms-full.txt b/docs/llms-full.txt index 086119307..78d2984d1 100644 --- a/docs/llms-full.txt +++ b/docs/llms-full.txt @@ -24606,7 +24606,7 @@ instead of generating them. | Format | Description | |--------|-------------| -| `{"ModelName": "package.Type"}` | Model-level: Skip generating `ModelName` and import from `package` | +| `{"ModelName": "package.Type"}` | Model-level: Skip generation; replace field and inheritance refs | | `{"Model.field": "package.Type"}` | Scoped: Override only specific field in specific model | **Common Use Cases:** @@ -24629,6 +24629,7 @@ instead of generating them. !!! note "Model-level overrides skip generation" When you specify a model-level override (without a dot in the key), the generator will **skip generating that model entirely** and import it from the specified package instead. + References to that model are replaced in field annotations and `allOf` inheritance base classes. ??? example "Examples" diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index 4bc92f937..7eddac1ef 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -1522,6 +1522,9 @@ def __init__( # noqa: PLR0912, PLR0915 self._type_override_imports: dict[str, Import] = { key: Import.from_full_path(value) for key, value in self.type_overrides.items() } + self._model_type_override_imports: dict[str, Import] = { + key: import_ for key, import_ in self._type_override_imports.items() if "." not in key + } self.read_only_write_only_model_type: ReadOnlyWriteOnlyModelType | None = config.read_only_write_only_model_type self.use_frozen_field: bool = config.use_frozen_field self.use_serialization_alias: bool = config.use_serialization_alias @@ -2805,14 +2808,12 @@ def __remove_overridden_models(self, models: list[DataModel]) -> list[DataModel] Only model-level overrides (keys without dots) cause model removal. Scoped overrides (ClassName.field) only affect specific fields. """ - if not self.type_overrides: + if not self._model_type_override_imports: return models - # Only model-level overrides (no dot) cause model removal - model_level_overrides = {k for k in self.type_overrides if "." not in k} - return [m for m in models if m.class_name not in model_level_overrides] + return [m for m in models if m.class_name not in self._model_type_override_imports] def __apply_type_overrides(self, models: list[DataModel]) -> None: - """Replace field type references with custom import types. + """Replace type references with custom import types. Supports two key formats: - Model-level: {"CustomType": "my_app.Type"} - applies to all references @@ -2823,14 +2824,31 @@ def __apply_type_overrides(self, models: list[DataModel]) -> None: if not self._type_override_imports: return for model in models: + fields_overridden = False for field in model.fields: # Check scoped override first: "ClassName.field_name" scoped_key = f"{model.class_name}.{field.name}" - if scoped_key in self._type_override_imports: - self._apply_override_to_field(field, self._type_override_imports[scoped_key]) - else: - # Apply model-level overrides to nested types - self._apply_override_to_data_type(field.data_type) + match self._type_override_imports.get(scoped_key): + case Import() as override_import: + self._apply_override_to_field(field, override_import) + fields_overridden = True + case None if self._model_type_override_imports: + # Apply model-level overrides to nested types + fields_overridden = self._apply_override_to_data_type(field.data_type) or fields_overridden + case None: + pass + if fields_overridden: + model.clear_imports_cache() + if not self._model_type_override_imports: + continue + base_classes_overridden = False + for base_class in model.base_classes: + if self._apply_override_to_data_type(base_class): + base_classes_overridden = True + for import_ in base_class.all_imports: + Parser._append_model_import(model, import_) + if base_classes_overridden: + model.clear_imports_cache() def _apply_override_to_field(self, field: DataModelFieldBase, override_import: Import) -> None: """Apply override to entire field's data_type.""" @@ -2841,16 +2859,27 @@ def _apply_override_to_field(self, field: DataModelFieldBase, override_import: I self.generation_store.detach_data_type_ref(field.data_type) self.generation_store.set_nested_data_types(field.data_type, []) - def _apply_override_to_data_type(self, data_type: DataType) -> None: + @staticmethod + def _append_model_import(model: DataModel, import_: Import) -> None: + """Append an import to a model once.""" + if import_ in model.imports: + return + model._additional_imports.append(import_) # noqa: SLF001 + + def _apply_override_to_data_type(self, data_type: DataType) -> bool: """Recursively apply model-level overrides to a DataType.""" - if data_type.reference and data_type.reference.name in self._type_override_imports: - override_import = self._type_override_imports[data_type.reference.name] + overridden = False + if data_type.reference and (override_import := self._model_type_override_imports.get(data_type.reference.name)): data_type.import_ = override_import data_type.alias = override_import.import_ self.generation_store.detach_data_type_ref(data_type) + overridden = True # Handle nested types (List[CustomType], Optional[CustomType], etc.) for nested in data_type.data_types: - self._apply_override_to_data_type(nested) + overridden = self._apply_override_to_data_type(nested) or overridden + if data_type.dict_key: + overridden = self._apply_override_to_data_type(data_type.dict_key) or overridden + return overridden @staticmethod def __disable_union_operator_for_forward_ref_parents(data_type: DataType) -> None: diff --git a/tests/data/expected/main/type_overrides_base_class.py b/tests/data/expected/main/type_overrides_base_class.py new file mode 100644 index 000000000..2b9f4de36 --- /dev/null +++ b/tests/data/expected/main/type_overrides_base_class.py @@ -0,0 +1,32 @@ +# generated by datamodel-codegen: +# filename: type_overrides_base_class.json +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from my_app.models import Base +from pydantic import BaseModel + + +class Other(BaseModel): + code: str | None = None + + +class Child(Base): + parent: Base | None = None + name: str | None = None + + +class BareChild(Base): + title: str | None = None + + +class Holder(BaseModel): + base: Base | None = None + child: Child | None = None + bare_child: BareChild | None = None + other: Other | None = None + + +class Model(BaseModel): + holder: Holder | None = None diff --git a/tests/data/expected/main/type_overrides_dict_key.py b/tests/data/expected/main/type_overrides_dict_key.py new file mode 100644 index 000000000..eb3693a7f --- /dev/null +++ b/tests/data/expected/main/type_overrides_dict_key.py @@ -0,0 +1,16 @@ +# generated by datamodel-codegen: +# filename: type_overrides_dict_key.json +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from my_app.keys import Key +from pydantic import BaseModel, RootModel + + +class ValueMap(RootModel[dict[Key, int]]): + root: dict[Key, int] + + +class Model(BaseModel): + values: ValueMap | None = None diff --git a/tests/data/jsonschema/type_overrides_base_class.json b/tests/data/jsonschema/type_overrides_base_class.json new file mode 100644 index 000000000..ccb2f30d0 --- /dev/null +++ b/tests/data/jsonschema/type_overrides_base_class.json @@ -0,0 +1,71 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "$defs": { + "Base": { + "type": "object", + "properties": { + "id": { + "type": "integer" + } + } + }, + "Other": { + "type": "object", + "properties": { + "code": { + "type": "string" + } + } + }, + "Child": { + "allOf": [ + { + "$ref": "#/$defs/Base" + } + ], + "properties": { + "parent": { + "$ref": "#/$defs/Base" + }, + "name": { + "type": "string" + } + } + }, + "BareChild": { + "allOf": [ + { + "$ref": "#/$defs/Base" + } + ], + "properties": { + "title": { + "type": "string" + } + } + }, + "Holder": { + "type": "object", + "properties": { + "base": { + "$ref": "#/$defs/Base" + }, + "child": { + "$ref": "#/$defs/Child" + }, + "bare_child": { + "$ref": "#/$defs/BareChild" + }, + "other": { + "$ref": "#/$defs/Other" + } + } + } + }, + "type": "object", + "properties": { + "holder": { + "$ref": "#/$defs/Holder" + } + } +} diff --git a/tests/data/jsonschema/type_overrides_dict_key.json b/tests/data/jsonschema/type_overrides_dict_key.json new file mode 100644 index 000000000..aa7846117 --- /dev/null +++ b/tests/data/jsonschema/type_overrides_dict_key.json @@ -0,0 +1,27 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "definitions": { + "Key": { + "type": "string", + "enum": [ + "primary", + "secondary" + ] + }, + "ValueMap": { + "type": "object", + "propertyNames": { + "$ref": "#/definitions/Key" + }, + "additionalProperties": { + "type": "integer" + } + } + }, + "type": "object", + "properties": { + "values": { + "$ref": "#/definitions/ValueMap" + } + } +} diff --git a/tests/main/test_main_general.py b/tests/main/test_main_general.py index 0eb16afda..4d7ac85ee 100644 --- a/tests/main/test_main_general.py +++ b/tests/main/test_main_general.py @@ -1501,12 +1501,13 @@ def test_dataclass_arguments_invalid(json_str: str, match: str) -> None: | Format | Description | |--------|-------------| -| `{"ModelName": "package.Type"}` | Model-level: Skip generating `ModelName` and import from `package` | +| `{"ModelName": "package.Type"}` | Model-level: Skip generation; replace field and inheritance refs | | `{"Model.field": "package.Type"}` | Scoped: Override only specific field in specific model | !!! note "Model-level overrides skip generation" When you specify a model-level override (without a dot in the key), the generator will **skip generating that model entirely** and import it from the specified package instead. + References to that model are replaced in field annotations and `allOf` inheritance base classes. **Common Use Cases:** @@ -1610,6 +1611,68 @@ def test_type_overrides_nested_types(output_file: Path) -> None: ) +@freeze_time(TIMESTAMP) +def test_type_overrides_model_level_base_class( + output_file: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Test model-level --type-overrides replaces base class references.""" + package_dir = tmp_path / "my_app" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + (package_dir / "models.py").write_text( + "from __future__ import annotations\n\n" + "from pydantic import BaseModel\n\n\n" + "class Base(BaseModel):\n" + " id: int | None = None\n", + encoding="utf-8", + ) + monkeypatch.syspath_prepend(str(tmp_path)) + + module_names = ("my_app", "my_app.models") + for module_name in module_names: + monkeypatch.delitem(sys.modules, module_name, raising=False) + + try: + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "type_overrides_base_class.json", + output_path=output_file, + input_file_type="jsonschema", + assert_func=assert_file_content, + expected_file="type_overrides_base_class.py", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--formatters", + "builtin", + "--type-overrides", + '{"Base": "my_app.models.Base"}', + ], + importable_module_name="generated_type_overrides_base_class", + importable_module_attribute="Holder", + ) + finally: + for module_name in module_names: + sys.modules.pop(module_name, None) + + +@freeze_time(TIMESTAMP) +def test_type_overrides_model_level_dict_key(output_file: Path) -> None: + """Test model-level --type-overrides replaces dict key references.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "type_overrides_dict_key.json", + output_path=output_file, + input_file_type="jsonschema", + assert_func=assert_file_content, + expected_file="type_overrides_dict_key.py", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--type-overrides", + '{"Key": "my_app.keys.Key"}', + ], + ) + + def test_skip_root_model(tmp_path: Path) -> None: """Test --skip-root-model flag functionality using generate().""" output_file = tmp_path / "output.py"