Fix aliased imports not applied to base classes and non-matching fields (#2981)

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -362,6 +362,29 @@ def iter_models_field_data_types(
                 yield model, field, data_type
 
 
+def _alias_base_class_imports(
+    model: DataModel,
+    aliased_imports: dict[tuple[str | None, str], Import],
+) -> None:
+    """Apply aliased imports to a model's base classes and their _additional_imports."""
+    for base_class in model.base_classes:
+        if not base_class.import_:
+            continue
+        key = (base_class.import_.from_, base_class.import_.import_)
+        if key not in aliased_imports:
+            continue
+        old_import = base_class.import_
+        aliased_import = aliased_imports[key]
+        base_class.type = aliased_import.alias  # type: ignore[assignment]
+        base_class.import_ = aliased_import
+        for i, additional_import in enumerate(model._additional_imports):  # pragma: no branch  # noqa: SLF001
+            if (
+                additional_import.from_ == old_import.from_ and additional_import.import_ == old_import.import_
+            ):  # pragma: no branch
+                model._additional_imports[i] = aliased_import  # noqa: SLF001
+                break
+
+
 ReferenceMapSet = dict[str, set[str]]
 SortedDataModels = dict[str, DataModel]
 
@@ -2388,21 +2411,31 @@ def __alias_shadowed_imports(  # noqa: PLR6301
         models: list[DataModel],
         all_model_field_names: set[str],
     ) -> None:
-        for _, model_field, data_type in iter_models_field_data_types(models):
-            if (
-                data_type
-                and data_type.import_
-                and data_type.type in all_model_field_names
-                and data_type.type == model_field.name
-            ):
-                alias = data_type.type + "_aliased"
-                data_type.type = alias
-                data_type.import_ = Import(
-                    from_=data_type.import_.from_,
-                    import_=data_type.import_.import_,
-                    alias=alias,
-                    reference_path=data_type.import_.reference_path,
-                )
+        aliased_imports: dict[tuple[str | None, str], Import] = {}
+        for _, _model_field, data_type in iter_models_field_data_types(models):
+            if data_type and data_type.import_ and data_type.type in all_model_field_names:
+                key = (data_type.import_.from_, data_type.import_.import_)
+                if key not in aliased_imports:
+                    aliased_imports[key] = Import(
+                        from_=data_type.import_.from_,
+                        import_=data_type.import_.import_,
+                        alias=data_type.type + "_aliased",
+                        reference_path=data_type.import_.reference_path,
+                    )
+
+        if not aliased_imports:
+            return
+
+        for _, _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  # type: ignore[assignment]
+                    data_type.import_ = aliased_import
+
+        for model in models:
+            _alias_base_class_imports(model, aliased_imports)
 
     def __apply_generic_base_class(  # noqa: PLR0912, PLR0914, PLR0915
         self,

tests/data/expected/main/openapi/shadowed_imports_base_and_fields.py

@@ -0,0 +1,23 @@
+# generated by datamodel-codegen:
+#   filename:  shadowed_imports_base_and_fields.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from mymodule.node import Node as Node_aliased
+from mymodule.other import Other
+from pydantic import BaseModel, ConfigDict
+
+
+class BaseItem(BaseModel):
+    id: str | None = None
+
+
+class MyModel(BaseItem):
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+    Node: Node_aliased | None = None
+    Node2: Node_aliased | None = None
+    other: Other | None = None
+    name: str | None = None

tests/data/expected/main/openapi/shadowed_imports_base_and_fields_custom_base.py

@@ -0,0 +1,23 @@
+# generated by datamodel-codegen:
+#   filename:  shadowed_imports_base_and_fields.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from mymodule.node import Node as Node_aliased
+from mymodule.other import Other
+from pydantic import ConfigDict
+
+
+class BaseItem(Node_aliased):
+    id: str | None = None
+
+
+class MyModel(BaseItem):
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+    Node: Node_aliased | None = None
+    Node2: Node_aliased | None = None
+    other: Other | None = None
+    name: str | None = None

tests/data/openapi/shadowed_imports_base_and_fields.yaml

@@ -0,0 +1,28 @@
+openapi: 3.0.0
+info:
+  title: Shadowed Imports Base and Fields Test
+  version: 0.0.1
+paths: {}
+components:
+  schemas:
+    BaseItem:
+      type: object
+      properties:
+        id:
+          type: string
+    MyModel:
+      allOf:
+        - $ref: '#/components/schemas/BaseItem'
+        - type: object
+          properties:
+            Node:
+              type: object
+              customTypePath: mymodule.node.Node
+            Node2:
+              type: object
+              customTypePath: mymodule.node.Node
+            other:
+              type: object
+              customTypePath: mymodule.other.Other
+            name:
+              type: string

tests/main/openapi/test_main_openapi.py

@@ -3712,6 +3712,30 @@ def test_main_openapi_shadowed_imports(output_file: Path) -> None:
     )
 
 
+def test_main_openapi_shadowed_imports_base_and_fields(output_file: Path) -> None:
+    """Test that aliased imports are applied to all fields, not just matching field names."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "shadowed_imports_base_and_fields.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file="shadowed_imports_base_and_fields.py",
+        extra_args=["--output-model-type", "pydantic_v2.BaseModel"],
+    )
+
+
+def test_main_openapi_shadowed_imports_base_and_fields_custom_base(output_file: Path) -> None:
+    """Test that aliased imports are applied to custom base classes."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "shadowed_imports_base_and_fields.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file="shadowed_imports_base_and_fields_custom_base.py",
+        extra_args=["--output-model-type", "pydantic_v2.BaseModel", "--base-class", "mymodule.node.Node"],
+    )
+
+
 def test_main_openapi_extra_fields_forbid(output_file: Path) -> None:
     """Test OpenAPI generation with extra fields forbidden."""
     run_main_and_assert(