Fix import handling in --collapse-root-models to exclude both Optional and Union (#2597)

* Fix handling of Optional imports in collapse-root-models functionality

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix import handling by using a list to exclude non-hashable imports in collapse-root-models

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -39,6 +39,7 @@
     IMPORT_ANNOTATIONS,
     IMPORT_LITERAL,
     IMPORT_OPTIONAL,
+    IMPORT_UNION,
     Import,
     Imports,
 )
@@ -1285,7 +1286,10 @@ def __collapse_root_models(  # noqa: PLR0912
                     original_field = get_most_of_parent(data_type, DataModelFieldBase)
                     if original_field:  # pragma: no cover
                         # TODO: Improve detection of reference type
-                        imports.append([i for i in original_field.imports if i != IMPORT_OPTIONAL])
+                        # Use list instead of set because Import is not hashable
+                        excluded_imports = [IMPORT_OPTIONAL, IMPORT_UNION]
+                        field_imports = [i for i in original_field.imports if i not in excluded_imports]
+                        imports.append(field_imports)
 
                     data_type.remove_reference()
 

tests/data/expected/main/jsonschema/jsonschema_collapse_root_models_with_optional.py

@@ -0,0 +1,13 @@
+# generated by datamodel-codegen:
+#   filename:  collapse_root_models_with_optional.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    field: Optional[str] = None

tests/data/jsonschema/collapse_root_models_with_optional.json

@@ -0,0 +1,20 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "type": "object",
+  "properties": {
+    "field": {
+      "anyOf": [
+        {"$ref": "#/$defs/StringType"},
+        {"$ref": "#/$defs/NullType"}
+      ]
+    }
+  },
+  "$defs": {
+    "StringType": {
+      "type": "string"
+    },
+    "NullType": {
+      "type": "null"
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2988,15 +2988,22 @@ def test_main_jsonschema_enum_object_values(output_file: Path) -> None:
 
 
 def test_main_jsonschema_collapse_root_models_empty_union(output_file: Path) -> None:
-    """Test that collapse-root-models with empty union fallback generates Any instead of invalid Union syntax.
-
-    This test covers the fix for issue #2161 where --collapse-root-models could generate
-    invalid Python syntax like Union[, str] when all union members are filtered out.
-    """
+    """Test that collapse-root-models with empty union fallback generates Any instead of invalid Union syntax."""
     run_main_and_assert(
         input_path=JSON_SCHEMA_DATA_PATH / "collapse_root_models_empty_union.json",
         output_path=output_file,
         input_file_type="jsonschema",
         assert_func=assert_file_content,
         extra_args=["--collapse-root-models"],
     )
+
+
+def test_main_jsonschema_collapse_root_models_with_optional(output_file: Path) -> None:
+    """Test that collapse-root-models correctly preserves Optional import when needed."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "collapse_root_models_with_optional.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        extra_args=["--collapse-root-models"],
+    )