Fix type overrides for base refs (#3589)

Author: koxudaxi

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"

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"

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:

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

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

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"
+    }
+  }
+}

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"
+    }
+  }
+}

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"