diff --git a/src/datamodel_code_generator/model/base.py b/src/datamodel_code_generator/model/base.py index a4a7d9868..db83ecc2b 100644 --- a/src/datamodel_code_generator/model/base.py +++ b/src/datamodel_code_generator/model/base.py @@ -678,6 +678,7 @@ def __init__( # noqa: PLR0913 self.reference.source = self + self.extra_template_data: dict[str, Any] if extra_template_data is not None: # The supplied defaultdict will either create a new entry, # or already contain a predefined entry for this type diff --git a/src/datamodel_code_generator/model/template/pydantic_v2/BaseModel.jinja2 b/src/datamodel_code_generator/model/template/pydantic_v2/BaseModel.jinja2 index d0e8c309d..828c4b416 100644 --- a/src/datamodel_code_generator/model/template/pydantic_v2/BaseModel.jinja2 +++ b/src/datamodel_code_generator/model/template/pydantic_v2/BaseModel.jinja2 @@ -7,7 +7,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme {{ description | escape_docstring | indent(4) }} """ {%- endif %} -{%- if not fields and not description and not config %} +{%- if not fields and not description and not config and not class_body_lines %} pass {%- endif %} {%- if config %} @@ -15,6 +15,9 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme {% include 'ConfigDict.jinja2' %} {%- endfilter %} {%- endif %} +{%- for line in class_body_lines %} + {{ line }} +{%- endfor %} {%- for field in fields %} {%- if not field.annotated and field.field %} {{ field.name }}: {{ field.type_hint }} = {{ field.field }} diff --git a/src/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2 b/src/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2 index 347d4dd36..af5de33b7 100644 --- a/src/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2 +++ b/src/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2 @@ -27,7 +27,10 @@ class {{ class_name }}({{ base_class }}{%- if fields -%}[{{get_type_hint(fields, {% include 'ConfigDict.jinja2' %} {%- endfilter %} {%- endif %} -{%- if not fields and not description %} +{%- for line in class_body_lines %} + {{ line }} +{%- endfor %} +{%- if not fields and not description and not config and not class_body_lines %} pass {%- else %} {%- set field = fields[0] %} diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index dc1bca6b5..0d1b4daaf 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -1548,6 +1548,32 @@ def __replace_unique_list_to_set(self, models: list[DataModel]) -> None: model_field.default = converted_default model_field.replace_data_type(set_data_type) + @classmethod + def __collect_set_item_references(cls, models: list[DataModel]) -> set[str]: + """Collect reference paths of all types used as set/frozenset items.""" + references: set[str] = set() + for model in models: + for field in model.fields: + for data_type in field.data_type.all_data_types: + if data_type.is_set or data_type.is_frozen_set: + for item_type in data_type.data_types: + references.update( + nested.reference.path for nested in item_type.all_data_types if nested.reference + ) + return references + + @classmethod + def __mark_set_item_models_hashable(cls, models: list[DataModel]) -> None: + """Mark models used as set/frozenset items with hash flag for __hash__ generation.""" + set_item_references = cls.__collect_set_item_references(models) + + for model in models: + if model.reference.path in set_item_references: + if isinstance(model, Enum): + continue + class_body_lines = model.extra_template_data.setdefault("class_body_lines", []) + class_body_lines.append("__hash__ = object.__hash__") + @classmethod def __set_reference_default_value_to_field(cls, models: list[DataModel]) -> None: for model in models: @@ -2965,6 +2991,8 @@ def _finalize_modules( module_to_import: dict[ModulePath, Imports], ) -> None: """Finalize module processing: apply generic base class and remove unused imports.""" + all_models = [model for ctx in contexts for model in ctx.models] + self.__mark_set_item_models_hashable(all_models) self.__apply_generic_base_class(contexts) for ctx in contexts: diff --git a/tests/data/expected/main/jsonschema/unique_items_enum_set.py b/tests/data/expected/main/jsonschema/unique_items_enum_set.py new file mode 100644 index 000000000..ac4b111a5 --- /dev/null +++ b/tests/data/expected/main/jsonschema/unique_items_enum_set.py @@ -0,0 +1,25 @@ +# generated by datamodel-codegen: +# filename: unique_items_enum_set.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel + + +class Status(Enum): + active = 'active' + inactive = 'inactive' + pending = 'pending' + + +class Item(BaseModel): + __hash__ = object.__hash__ + name: str | None = None + + +class Container(BaseModel): + statuses: set[Status] | None = None + items: set[Item] | None = None diff --git a/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_generic_container_types_set.py b/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_generic_container_types_set.py index 5616c0270..ddd77b3fc 100644 --- a/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_generic_container_types_set.py +++ b/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_generic_container_types_set.py @@ -10,6 +10,7 @@ class Pet(BaseModel): + __hash__ = object.__hash__ id: int = Field(..., ge=0, le=9223372036854775807) name: str = Field(..., max_length=256) tag: str | None = Field(None, max_length=64) diff --git a/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_standard_collections_set.py b/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_standard_collections_set.py index 908d73f33..383963678 100644 --- a/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_standard_collections_set.py +++ b/tests/data/expected/main/openapi/with_field_constraints_pydantic_v2_use_standard_collections_set.py @@ -8,6 +8,7 @@ class Pet(BaseModel): + __hash__ = object.__hash__ id: int = Field(..., ge=0, le=9223372036854775807) name: str = Field(..., max_length=256) tag: str | None = Field(None, max_length=64) diff --git a/tests/data/jsonschema/unique_items_enum_set.json b/tests/data/jsonschema/unique_items_enum_set.json new file mode 100644 index 000000000..debe60bba --- /dev/null +++ b/tests/data/jsonschema/unique_items_enum_set.json @@ -0,0 +1,29 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "Container", + "type": "object", + "definitions": { + "Status": { + "type": "string", + "enum": ["active", "inactive", "pending"] + }, + "Item": { + "type": "object", + "properties": { + "name": {"type": "string"} + } + } + }, + "properties": { + "statuses": { + "type": "array", + "uniqueItems": true, + "items": {"$ref": "#/definitions/Status"} + }, + "items": { + "type": "array", + "uniqueItems": true, + "items": {"$ref": "#/definitions/Item"} + } + } +} diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py index 6dda2ee8b..cf7a0e8c9 100644 --- a/tests/main/jsonschema/test_main_jsonschema.py +++ b/tests/main/jsonschema/test_main_jsonschema.py @@ -7903,3 +7903,21 @@ def test_validators_requires_pydantic_v2(output_file: Path, tmp_path: Path, caps capsys=capsys, expected_stderr_contains="--validators option requires Pydantic v2", ) + + +@PYDANTIC_V2_SKIP +def test_unique_items_enum_set(output_file: Path) -> None: + """Test set with enum items does not add __hash__ to enum (already hashable).""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "unique_items_enum_set.json", + output_path=output_file, + input_file_type="jsonschema", + assert_func=assert_file_content, + expected_file="unique_items_enum_set.py", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--use-unique-items-as-set", + "--use-standard-collections", + ], + )