From 935384a76da900cbe10c38f58582b6921025c324 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Fri, 19 Dec 2025 19:38:53 +0000 Subject: [PATCH 1/3] fix: Enhance deep copy functionality and improve import management in core classes --- src/datamodel_code_generator/imports.py | 39 ++++++++---- src/datamodel_code_generator/model/base.py | 6 ++ .../model/dataclass.py | 21 ++++++- src/datamodel_code_generator/types.py | 6 +- .../__init__.py | 3 + .../schema_a.py | 20 +++++++ .../schema_b.py | 15 +++++ .../shared.py | 14 +++++ tests/main/jsonschema/test_main_jsonschema.py | 18 ++++++ tests/test_imports.py | 59 +++++++++++++++++++ tests/test_types.py | 19 ++++++ 11 files changed, 202 insertions(+), 18 deletions(-) create mode 100644 tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/__init__.py create mode 100644 tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_a.py create mode 100644 tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/schema_b.py create mode 100644 tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/shared.py diff --git a/src/datamodel_code_generator/imports.py b/src/datamodel_code_generator/imports.py index 370610261..0ae4e43ba 100644 --- 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 index 36654166f..75de4d497 100644 --- 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 index 79534201b..fa9a3853f 100644 --- a/src/datamodel_code_generator/model/dataclass.py +++ b/src/datamodel_code_generator/model/dataclass.py @@ -21,6 +21,7 @@ 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: @@ -28,8 +29,6 @@ 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 index 9f0ee0694..21c4c399e 100644 --- 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/reuse_scope_tree_dataclass_frozen/__init__.py b/tests/data/expected/main/jsonschema/reuse_scope_tree_dataclass_frozen/__init__.py new file mode 100644 index 000000000..14bcc7af7 --- /dev/null +++ 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 new file mode 100644 index 000000000..6fbcce447 --- /dev/null +++ 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 new file mode 100644 index 000000000..5f133587c --- /dev/null +++ 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 new file mode 100644 index 000000000..ffbd85ac9 --- /dev/null +++ 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/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py index b3124754c..d6779f139 100644 --- a/tests/main/jsonschema/test_main_jsonschema.py +++ b/tests/main/jsonschema/test_main_jsonschema.py @@ -3923,6 +3923,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/test_imports.py b/tests/test_imports.py index 34864e0b4..f9bd00224 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -96,3 +96,62 @@ def test_extract_future_with_alias() -> None: assert "annotations as ann" in str(future) assert "__future__" not in imports assert "__future__" not in imports.alias + + +def test_remove_nonexistent_import() -> None: + """Test that removing non-existent import doesn't crash.""" + imports = Imports() + imports.append(Import(from_="typing", import_="Optional")) + + imports.remove(Import(from_="typing", import_="List")) + + assert str(imports) == "from typing import Optional" + + +def test_remove_double_removal() -> None: + """Test that double removal of same import doesn't crash.""" + imports = Imports() + imports.append(Import(from_="typing", import_="Optional")) + + imports.remove(Import(from_="typing", import_="Optional")) + imports.remove(Import(from_="typing", import_="Optional")) + + assert not str(imports) + + +def test_remove_cleans_up_counter() -> None: + """Test that remove() properly cleans up counter entries.""" + imports = Imports() + imports.append(Import(from_="typing", import_="Optional")) + + assert imports.counter.get(("typing", "Optional")) == 1 + + imports.remove(Import(from_="typing", import_="Optional")) + + assert ("typing", "Optional") not in imports.counter + + +def test_remove_cleans_up_reference_paths() -> None: + """Test that remove() properly cleans up reference_paths.""" + imports = Imports() + imports.append(Import(from_="typing", import_="Optional", reference_path="/test/path")) + + assert "/test/path" in imports.reference_paths + + imports.remove(Import(from_="typing", import_="Optional", reference_path="/test/path")) + + assert "/test/path" not in imports.reference_paths + + +def test_extract_future_moves_reference_paths() -> None: + """Test that extract_future() moves reference_paths for __future__ imports.""" + imports = Imports() + imports.append(Import(from_="__future__", import_="annotations", reference_path="/future/ref")) + imports.append(Import(from_="typing", import_="Optional", reference_path="/typing/ref")) + + future = imports.extract_future() + + assert "/future/ref" in future.reference_paths + assert "/future/ref" not in imports.reference_paths + assert "/typing/ref" in imports.reference_paths + assert "/typing/ref" not in future.reference_paths diff --git a/tests/test_types.py b/tests/test_types.py index 775c63c9a..6d32e31f1 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -133,3 +133,22 @@ def test_get_optional_type(input_: str, use_union_operator: bool, expected: str) def test_remove_none_from_union(type_str: str, use_union_operator: bool, expected: str) -> None: """Test _remove_none_from_union function with various type strings.""" assert _remove_none_from_union(type_str, use_union_operator=use_union_operator) == expected + + +@pytest.mark.parametrize( + ("type_str", "use_union_operator", "expected"), + [ + ("(", False, "("), + (")", False, ")"), + ("()", False, "()"), + ("a(", False, "a("), + ("constr()", False, "constr()"), + ("constr(pattern=')')", False, "constr(pattern=')')"), + ("Union[constr()]", False, "constr()"), + ("a | b", True, "a | b"), + ("(a)", True, "(a)"), + ], +) +def test_remove_none_from_union_short_strings(type_str: str, use_union_operator: bool, expected: str) -> None: + """Test _remove_none_from_union with short strings to verify index bounds safety.""" + assert _remove_none_from_union(type_str, use_union_operator=use_union_operator) == expected From 8f8a4de677defdfb4f7b3136355e22f68dbcfa0e Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Fri, 19 Dec 2025 19:43:52 +0000 Subject: [PATCH 2/3] fix: Import Reference in dataclass.py to resolve missing dependency issues --- src/datamodel_code_generator/model/dataclass.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/datamodel_code_generator/model/dataclass.py b/src/datamodel_code_generator/model/dataclass.py index fa9a3853f..bdf947a9a 100644 --- a/src/datamodel_code_generator/model/dataclass.py +++ b/src/datamodel_code_generator/model/dataclass.py @@ -21,7 +21,6 @@ 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: @@ -29,6 +28,8 @@ 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.""" @@ -92,6 +93,8 @@ def __init__( # noqa: PLR0913 def create_reuse_model(self, base_ref: Reference) -> DataClass: """Create inherited model with empty fields pointing to base reference.""" + from datamodel_code_generator.reference import Reference # noqa: PLC0415 + return self.__class__( fields=[], base_classes=[base_ref], From 1c7df5440e52a170f2deea48ccf9c55447a41de1 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Fri, 19 Dec 2025 19:56:03 +0000 Subject: [PATCH 3/3] fix: Resolve import issues and enhance deep copy functionality in core classes --- .../model/dataclass.py | 5 +- .../copy_deep_pattern_properties.py | 38 ++++++++++++++ .../copy_deep_pattern_properties.json | 50 +++++++++++++++++++ tests/main/jsonschema/test_main_jsonschema.py | 17 +++++++ tests/model/test_base.py | 30 +++++++++++ 5 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tests/data/expected/main/jsonschema/copy_deep_pattern_properties.py create mode 100644 tests/data/jsonschema/copy_deep_pattern_properties.json diff --git a/src/datamodel_code_generator/model/dataclass.py b/src/datamodel_code_generator/model/dataclass.py index bdf947a9a..fa9a3853f 100644 --- a/src/datamodel_code_generator/model/dataclass.py +++ b/src/datamodel_code_generator/model/dataclass.py @@ -21,6 +21,7 @@ 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: @@ -28,8 +29,6 @@ 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.""" @@ -93,8 +92,6 @@ def __init__( # noqa: PLR0913 def create_reuse_model(self, base_ref: Reference) -> DataClass: """Create inherited model with empty fields pointing to base reference.""" - from datamodel_code_generator.reference import Reference # noqa: PLC0415 - return self.__class__( fields=[], base_classes=[base_ref], diff --git a/tests/data/expected/main/jsonschema/copy_deep_pattern_properties.py b/tests/data/expected/main/jsonschema/copy_deep_pattern_properties.py new file mode 100644 index 000000000..27e1b691f --- /dev/null +++ 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/jsonschema/copy_deep_pattern_properties.json b/tests/data/jsonschema/copy_deep_pattern_properties.json new file mode 100644 index 000000000..dfe699715 --- /dev/null +++ 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 index d6779f139..1e6e4513b 100644 --- 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( diff --git a/tests/model/test_base.py b/tests/model/test_base.py index 80cf3dfbc..6b98a38c9 100644 --- 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"