fix: retain imports with aliases during cleanup process (#2569)

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -1341,6 +1341,39 @@ def __alias_shadowed_imports(  # noqa: PLR6301
                             reference_path=model_field.data_type.import_.reference_path,
                         )
 
+    @classmethod
+    def _collect_used_names_from_models(cls, models: list[DataModel]) -> set[str]:
+        """Collect identifiers referenced by models before rendering."""
+        names: set[str] = set()
+
+        def add(name: str | None) -> None:
+            if not name:
+                return
+            # first segment is sufficient to match import target or alias
+            names.add(name.split(".")[0])
+
+        def walk_data_type(data_type: DataType) -> None:
+            add(data_type.alias or data_type.type)
+            if data_type.reference:
+                add(data_type.reference.short_name)
+            for child in data_type.data_types:
+                walk_data_type(child)
+            if data_type.dict_key:
+                walk_data_type(data_type.dict_key)
+
+        for model in models:
+            add(model.class_name)
+            add(model.duplicate_class_name)
+            for base in model.base_classes:
+                add(base.type_hint)
+            for import_ in model.imports:
+                add(import_.alias or import_.import_.split(".")[-1])
+            for field in model.fields:
+                add(field.name)
+                add(field.alias)
+                walk_data_type(field.data_type)
+        return names
+
     def parse(  # noqa: PLR0912, PLR0914, PLR0915
         self,
         with_import: bool | None = True,  # noqa: FBT001, FBT002
@@ -1466,12 +1499,14 @@ class Processed(NamedTuple):
 
         for processed_model in processed_models:
             # postprocess imports to remove unused imports.
-            model_code = str("\n".join([str(m) for m in processed_model.models]))
+            used_names = self._collect_used_names_from_models(processed_model.models)
             unused_imports = [
                 (from_, import_)
                 for from_, imports_ in processed_model.imports.items()
                 for import_ in imports_
-                if import_ not in model_code
+                if not {processed_model.imports.alias.get(from_, {}).get(import_, import_), import_}.intersection(
+                    used_names
+                )
             ]
             for from_, import_ in unused_imports:
                 processed_model.imports.remove(Import(from_=from_, import_=import_))

tests/data/expected/main/jsonschema/alias_import_alias/__init__.py

@@ -0,0 +1,3 @@
+# generated by datamodel-codegen:
+#   filename:  alias_import_alias
+#   timestamp: 2019-07-26T00:00:00+00:00

tests/data/expected/main/jsonschema/alias_import_alias/a_schema.py

@@ -0,0 +1,11 @@
+# generated by datamodel-codegen:
+#   filename:  a.schema.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class TypeA(BaseModel):
+    value: str

tests/data/expected/main/jsonschema/alias_import_alias/b_schema.py

@@ -0,0 +1,13 @@
+# generated by datamodel-codegen:
+#   filename:  b.schema.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from . import a_schema as a_schema_1
+
+
+class Container(BaseModel):
+    a_schema: a_schema_1.TypeA

tests/data/jsonschema/alias_import_alias/a.schema.json

@@ -0,0 +1,11 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "TypeA",
+  "type": "object",
+  "properties": {
+    "value": {
+      "type": "string"
+    }
+  },
+  "required": ["value"]
+}

tests/data/jsonschema/alias_import_alias/b.schema.json

@@ -0,0 +1,11 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "Container",
+  "type": "object",
+  "properties": {
+    "a_schema": {
+      "$ref": "a.schema.json"
+    }
+  },
+  "required": ["a_schema"]
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2405,6 +2405,15 @@ def test_main_invalid_import_name(output_dir: Path) -> None:
     )
 
 
+def test_main_alias_import_alias(output_dir: Path) -> None:
+    """Ensure imports with aliases are retained after cleanup."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "alias_import_alias",
+        output_path=output_dir,
+        expected_directory=EXPECTED_JSON_SCHEMA_PATH / "alias_import_alias",
+    )
+
+
 @pytest.mark.parametrize(
     ("output_model", "expected_output"),
     [