fix: Improve state management in Imports, DataType, and DataModel classes (#2705)

koxudaxi · web-flow · commit 1b5d3fb3bf69 · 2025-12-20T05:01:32.000+09:00
* fix: Enhance deep copy functionality and improve import management in core classes

* fix: Import Reference in dataclass.py to resolve missing dependency issues

* fix: Resolve import issues and enhance deep copy functionality in core classes
diff --git a/src/datamodel_code_generator/imports.py b/src/datamodel_code_generator/imports.py
@@ -88,27 +88,39 @@ def append(self, imports: Import | Iterable[Import] | None) -> None:
                     if import_.alias:
                         self.alias[import_.from_][import_.import_] = import_.alias
 
-    def remove(self, imports: Import | Iterable[Import]) -> None:
+    def remove(self, imports: Import | Iterable[Import]) -> None:  # noqa: PLR0912
         """Remove one or more imports from the collection."""
         if isinstance(imports, Import):  # pragma: no cover
             imports = [imports]
         for import_ in imports:
             if "." in import_.import_:  # pragma: no cover
-                self.counter[None, import_.import_] -= 1
-                if self.counter[None, import_.import_] == 0:  # pragma: no cover
-                    self[None].remove(import_.import_)
-                    if not self[None]:
-                        del self[None]
+                key = (None, import_.import_)
+                if self.counter.get(key, 0) <= 0:
+                    continue
+                self.counter[key] -= 1
+                if self.counter[key] == 0:  # pragma: no cover
+                    del self.counter[key]
+                    if None in self and import_.import_ in self[None]:
+                        self[None].remove(import_.import_)
+                        if not self[None]:
+                            del self[None]
             else:
-                self.counter[import_.from_, import_.import_] -= 1  # pragma: no cover
-                if self.counter[import_.from_, import_.import_] == 0:  # pragma: no cover
-                    self[import_.from_].remove(import_.import_)
-                    if not self[import_.from_]:
-                        del self[import_.from_]
-                    if import_.alias:  # pragma: no cover
+                key = (import_.from_, import_.import_)
+                if self.counter.get(key, 0) <= 0:
+                    continue
+                self.counter[key] -= 1  # pragma: no cover
+                if self.counter[key] == 0:  # pragma: no cover
+                    del self.counter[key]
+                    if import_.from_ in self and import_.import_ in self[import_.from_]:
+                        self[import_.from_].remove(import_.import_)
+                        if not self[import_.from_]:
+                            del self[import_.from_]
+                    if import_.alias and import_.from_ in self.alias and import_.import_ in self.alias[import_.from_]:
                         del self.alias[import_.from_][import_.import_]
                         if not self.alias[import_.from_]:
                             del self.alias[import_.from_]
+            if import_.reference_path and import_.reference_path in self.reference_paths:
+                del self.reference_paths[import_.reference_path]
 
     def remove_referenced_imports(self, reference_path: str) -> None:
         """Remove imports associated with a reference path."""
@@ -126,6 +138,9 @@ def extract_future(self) -> Imports:
                     future.counter[key] = self.counter.pop(key)
             if future_key in self.alias:
                 future.alias[future_key] = self.alias.pop(future_key)
+            for ref_path, import_ in list(self.reference_paths.items()):
+                if import_.from_ == future_key:
+                    future.reference_paths[ref_path] = self.reference_paths.pop(ref_path)
         return future
 
     def add_export(self, name: str) -> None:
diff --git a/src/datamodel_code_generator/model/base.py b/src/datamodel_code_generator/model/base.py
@@ -328,9 +328,12 @@ def copy_deep(self) -> Self:
         """Create a deep copy of this field to avoid mutating the original."""
         copied = self.copy()
         copied.parent = None
+        copied.extras = deepcopy(self.extras)
         copied.data_type = self.data_type.copy()
         if self.data_type.data_types:
             copied.data_type.data_types = [dt.copy() for dt in self.data_type.data_types]
+        if self.data_type.dict_key:
+            copied.data_type.dict_key = self.data_type.dict_key.copy()
         return copied
 
     def replace_data_type(self, new_data_type: DataType, *, clear_old_parent: bool = True) -> None:
@@ -549,6 +552,9 @@ def create_reuse_model(self, base_ref: Reference) -> Self:
                 path=self.reference.path + "/reuse",
             ),
             custom_template_dir=self._custom_template_dir,
+            custom_base_class=self.custom_base_class,
+            keyword_only=self.keyword_only,
+            treat_dot_as_module=self._treat_dot_as_module,
         )
 
     def replace_children_in_models(self, models: list[DataModel], new_ref: Reference) -> None:
diff --git a/src/datamodel_code_generator/model/dataclass.py b/src/datamodel_code_generator/model/dataclass.py
@@ -21,15 +21,14 @@
 from datamodel_code_generator.model.pydantic.base_model import Constraints  # noqa: TC001 # needed for pydantic
 from datamodel_code_generator.model.types import DataTypeManager as _DataTypeManager
 from datamodel_code_generator.model.types import type_map_factory
+from datamodel_code_generator.reference import Reference
 from datamodel_code_generator.types import DataType, StrictTypes, Types, chain_as_tuple
 
 if TYPE_CHECKING:
     from collections import defaultdict
     from collections.abc import Sequence
     from pathlib import Path
 
-    from datamodel_code_generator.reference import Reference
-
 
 def has_field_assignment(field: DataModelFieldBase) -> bool:
     """Check if a dataclass field has a default value or field() assignment."""
@@ -91,6 +90,24 @@ def __init__(  # noqa: PLR0913
             if keyword_only:
                 self.dataclass_arguments["kw_only"] = True
 
+    def create_reuse_model(self, base_ref: Reference) -> DataClass:
+        """Create inherited model with empty fields pointing to base reference."""
+        return self.__class__(
+            fields=[],
+            base_classes=[base_ref],
+            description=self.description,
+            reference=Reference(
+                name=self.name,
+                path=self.reference.path + "/reuse",
+            ),
+            custom_template_dir=self._custom_template_dir,
+            custom_base_class=self.custom_base_class,
+            keyword_only=self.keyword_only,
+            frozen=self.frozen,
+            treat_dot_as_module=self._treat_dot_as_module,
+            dataclass_arguments=self.dataclass_arguments,
+        )
+
 
 class DataModelField(DataModelFieldBase):
     """Field implementation for dataclass models."""
diff --git a/src/datamodel_code_generator/types.py b/src/datamodel_code_generator/types.py
@@ -207,12 +207,10 @@ def _remove_none_from_union(type_: str, *, use_union_operator: bool) -> str:  #
         elif char == "]" and in_constr == 0:
             inner_count -= 1
         elif char == "(":
-            if current_part.strip().startswith("constr(") and current_part[-2] != "\\":
-                # non-escaped opening round bracket found inside constraint string expression
+            if current_part.strip().startswith("constr(") and (len(current_part) < 2 or current_part[-2] != "\\"):  # noqa: PLR2004
                 in_constr += 1
         elif char == ")":
-            if in_constr > 0 and current_part[-2] != "\\":
-                # non-escaped closing round bracket found inside constraint string expression
+            if in_constr > 0 and (len(current_part) < 2 or current_part[-2] != "\\"):  # noqa: PLR2004
                 in_constr -= 1
         elif char == separator and inner_count == 0 and in_constr == 0:
             part = current_part[:-1].strip()
diff --git a/tests/data/expected/main/jsonschema/copy_deep_pattern_properties.py b/tests/data/expected/main/jsonschema/copy_deep_pattern_properties.py
@@ -0,0 +1,38 @@
+# generated by datamodel-codegen:
+#   filename:  copy_deep_pattern_properties.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import AwareDatetime, BaseModel, Field, constr
+
+
+class MetadataRequest(BaseModel):
+    tags: Optional[dict[constr(pattern=r'^[a-z][a-z0-9_]*$'), str]] = Field(
+        None, description='Dynamic key-value metadata'
+    )
+
+
+class Metadata(BaseModel):
+    tags: Optional[dict[constr(pattern=r'^[a-z][a-z0-9_]*$'), str]] = Field(
+        None, description='Dynamic key-value metadata'
+    )
+    created_at: Optional[AwareDatetime] = None
+
+
+class ExtendedMetadataRequest(BaseModel):
+    tags: Optional[dict[constr(pattern=r'^[a-z][a-z0-9_]*$'), str]] = Field(
+        None, description='Dynamic key-value metadata'
+    )
+    owner: Optional[str] = None
+
+
+class ExtendedMetadata(Metadata):
+    id: int
+    owner: Optional[str] = None
+
+
+class Model(BaseModel):
+    metadata: Optional[ExtendedMetadata] = None
diff --git a/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/__init__.py b/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/__init__.py
@@ -0,0 +1,3 @@
+# generated by datamodel-codegen:
+#   filename:  reuse_scope_tree_dataclass
+#   timestamp: 2019-07-26T00:00:00+00:00
diff --git a/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_a.py b/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_a.py
@@ -0,0 +1,20 @@
+# generated by datamodel-codegen:
+#   filename:  reuse_scope_tree_dataclass
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from .shared import SharedModel as SharedModel_1
+
+
+@dataclass(frozen=True)
+class SharedModel(SharedModel_1):
+    pass
+
+
+@dataclass(frozen=True)
+class Model:
+    data: Optional[SharedModel] = None
diff --git a/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_b.py b/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_b.py
@@ -0,0 +1,15 @@
+# generated by datamodel-codegen:
+#   filename:  schema_b.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from . import shared
+
+
+@dataclass(frozen=True)
+class Model:
+    info: Optional[shared.SharedModel] = None
diff --git a/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/shared.py b/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/shared.py
@@ -0,0 +1,14 @@
+# generated by datamodel-codegen:
+#   filename:  shared.py
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class SharedModel:
+    id: Optional[int] = None
+    name: Optional[str] = None
diff --git a/tests/data/jsonschema/copy_deep_pattern_properties.json b/tests/data/jsonschema/copy_deep_pattern_properties.json
@@ -0,0 +1,50 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "definitions": {
+    "Metadata": {
+      "type": "object",
+      "properties": {
+        "tags": {
+          "type": "object",
+          "description": "Dynamic key-value metadata",
+          "patternProperties": {
+            "^[a-z][a-z0-9_]*$": {
+              "type": "string"
+            }
+          }
+        },
+        "created_at": {
+          "type": "string",
+          "format": "date-time",
+          "readOnly": true
+        }
+      }
+    },
+    "ExtendedMetadata": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/Metadata"
+        },
+        {
+          "type": "object",
+          "required": ["id"],
+          "properties": {
+            "id": {
+              "type": "integer",
+              "readOnly": true
+            },
+            "owner": {
+              "type": "string"
+            }
+          }
+        }
+      ]
+    }
+  },
+  "type": "object",
+  "properties": {
+    "metadata": {
+      "$ref": "#/definitions/ExtendedMetadata"
+    }
+  }
+}
diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py
@@ -2652,6 +2652,23 @@ def test_main_jsonschema_pattern_properties_by_reference(output_file: Path) -> N
     )
 
 
+def test_main_jsonschema_copy_deep_pattern_properties(output_file: Path) -> None:
+    """Test copy_deep properly preserves dict_key from patternProperties during allOf inheritance."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "copy_deep_pattern_properties.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        expected_file="copy_deep_pattern_properties.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--read-only-write-only-model-type",
+            "all",
+        ],
+    )
+
+
 def test_main_dataclass_field(output_file: Path) -> None:
     """Test dataclass field generation."""
     run_main_and_assert(
@@ -3923,6 +3940,24 @@ def test_main_jsonschema_reuse_scope_tree_dataclass(output_dir: Path) -> None:
     )
 
 
+def test_main_jsonschema_reuse_scope_tree_dataclass_frozen(output_dir: Path) -> None:
+    """Test --reuse-scope=tree with frozen dataclasses preserves frozen in inherited models."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "reuse_scope_tree_dataclass",
+        output_path=output_dir,
+        expected_directory=EXPECTED_JSON_SCHEMA_PATH / "reuse_scope_tree_dataclass_frozen",
+        input_file_type="jsonschema",
+        extra_args=[
+            "--reuse-model",
+            "--reuse-scope",
+            "tree",
+            "--output-model-type",
+            "dataclasses.dataclass",
+            "--frozen",
+        ],
+    )
+
+
 def test_main_jsonschema_reuse_scope_tree_typeddict(output_dir: Path) -> None:
     """Test --reuse-scope=tree with TypedDict output type (no inheritance, direct reference)."""
     run_main_and_assert(
diff --git a/tests/model/test_base.py b/tests/model/test_base.py
@@ -338,3 +338,33 @@ def test_get_module_path_without_file_path_parametrized(
     """Test module path generation without file path for various module names."""
     result = get_module_path(name, None, treat_dot_as_module=treat_dot_as_module)
     assert result == expected
+
+
+def test_copy_deep_with_dict_key() -> None:
+    """Test that copy_deep properly copies dict_key."""
+    dict_key_type = DataType(type="str")
+    data_type = DataType(is_dict=True, dict_key=dict_key_type)
+    field = DataModelFieldBase(name="a", data_type=data_type, required=True)
+
+    copied = field.copy_deep()
+
+    assert copied.data_type.dict_key is not None
+    assert copied.data_type.dict_key is not field.data_type.dict_key
+    assert copied.data_type.dict_key.type == "str"
+
+
+def test_copy_deep_with_extras() -> None:
+    """Test that copy_deep properly deep copies extras."""
+    field = DataModelFieldBase(
+        name="a",
+        data_type=DataType(type="str"),
+        required=True,
+        extras={"key": "value", "nested": {"inner": 1}},
+    )
+
+    copied = field.copy_deep()
+
+    assert copied.extras is not field.extras
+    assert copied.extras == {"key": "value", "nested": {"inner": 1}}
+    copied.extras["key"] = "modified"
+    assert field.extras["key"] == "value"
diff --git a/tests/test_imports.py b/tests/test_imports.py
diff --git a/tests/test_types.py b/tests/test_types.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# generated by datamodel-codegen:`
	`2`	`+# filename: reuse_scope_tree_dataclass`
	`3`	`+# timestamp: 2019-07-26T00:00:00+00:00`