datamodel-code-generator 0.11.12__py3-none-any.whl → 0.45.0__py3-none-any.whl

datamodel_code_generator/model/dataclass.py

@@ -0,0 +1,211 @@
+"""Python dataclass model generator.
+
+Generates Python dataclasses using the @dataclass decorator.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
+
+from datamodel_code_generator import DataclassArguments, DatetimeClassType, PythonVersion, PythonVersionMin
+from datamodel_code_generator.imports import (
+    IMPORT_DATE,
+    IMPORT_DATETIME,
+    IMPORT_TIME,
+    IMPORT_TIMEDELTA,
+    Import,
+)
+from datamodel_code_generator.model import DataModel, DataModelFieldBase
+from datamodel_code_generator.model.base import UNDEFINED
+from datamodel_code_generator.model.imports import IMPORT_DATACLASS, IMPORT_FIELD
+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.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."""
+    return bool(field.field) or not (
+        field.required or (field.represented_default == "None" and field.strip_default_none)
+    )
+
+
+class DataClass(DataModel):
+    """DataModel implementation for Python dataclasses."""
+
+    TEMPLATE_FILE_PATH: ClassVar[str] = "dataclass.jinja2"
+    DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = (IMPORT_DATACLASS,)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        reference: Reference,
+        fields: list[DataModelFieldBase],
+        decorators: list[str] | None = None,
+        base_classes: list[Reference] | None = None,
+        custom_base_class: str | None = None,
+        custom_template_dir: Path | None = None,
+        extra_template_data: defaultdict[str, dict[str, Any]] | None = None,
+        methods: list[str] | None = None,
+        path: Path | None = None,
+        description: str | None = None,
+        default: Any = UNDEFINED,
+        nullable: bool = False,
+        keyword_only: bool = False,
+        frozen: bool = False,
+        treat_dot_as_module: bool = False,
+        dataclass_arguments: DataclassArguments | None = None,
+    ) -> None:
+        """Initialize dataclass with fields sorted by field assignment requirement."""
+        super().__init__(
+            reference=reference,
+            fields=sorted(fields, key=has_field_assignment),
+            decorators=decorators,
+            base_classes=base_classes,
+            custom_base_class=custom_base_class,
+            custom_template_dir=custom_template_dir,
+            extra_template_data=extra_template_data,
+            methods=methods,
+            path=path,
+            description=description,
+            default=default,
+            nullable=nullable,
+            keyword_only=keyword_only,
+            frozen=frozen,
+            treat_dot_as_module=treat_dot_as_module,
+        )
+        if dataclass_arguments is not None:
+            self.dataclass_arguments = dataclass_arguments
+        else:
+            self.dataclass_arguments = {}
+            if frozen:
+                self.dataclass_arguments["frozen"] = True
+            if keyword_only:
+                self.dataclass_arguments["kw_only"] = True
+
+
+class DataModelField(DataModelFieldBase):
+    """Field implementation for dataclass models."""
+
+    _FIELD_KEYS: ClassVar[set[str]] = {
+        "default_factory",
+        "init",
+        "repr",
+        "hash",
+        "compare",
+        "metadata",
+        "kw_only",
+    }
+    constraints: Optional[Constraints] = None  # noqa: UP045
+
+    def process_const(self) -> None:
+        """Process const field constraint using literal type."""
+        self._process_const_as_literal()
+
+    @property
+    def imports(self) -> tuple[Import, ...]:
+        """Get imports including field() if needed."""
+        field = self.field
+        if field and field.startswith("field("):
+            return chain_as_tuple(super().imports, (IMPORT_FIELD,))
+        return super().imports
+
+    @property
+    def field(self) -> str | None:
+        """For backwards compatibility."""
+        result = str(self)
+        if not result:
+            return None
+        return result
+
+    def __str__(self) -> str:
+        """Generate field() call or default value representation."""
+        data: dict[str, Any] = {k: v for k, v in self.extras.items() if k in self._FIELD_KEYS}
+
+        if self.default != UNDEFINED and self.default is not None:
+            data["default"] = self.default
+
+        if self.required:
+            data = {
+                k: v
+                for k, v in data.items()
+                if k
+                not in {
+                    "default",
+                    "default_factory",
+                }
+            }
+
+        if not data:
+            return ""
+
+        if len(data) == 1 and "default" in data:
+            default = data["default"]
+
+            if isinstance(default, (list, dict, set)):
+                if default:
+                    from datamodel_code_generator.model.base import repr_set_sorted  # noqa: PLC0415
+
+                    default_repr = repr_set_sorted(default) if isinstance(default, set) else repr(default)
+                    return f"field(default_factory=lambda: {default_repr})"
+                return f"field(default_factory={type(default).__name__})"
+            return repr(default)
+        kwargs = [f"{k}={v if k == 'default_factory' else repr(v)}" for k, v in data.items()]
+        return f"field({', '.join(kwargs)})"
+
+
+class DataTypeManager(_DataTypeManager):
+    """Type manager for dataclass models."""
+
+    def __init__(  # noqa: PLR0913, PLR0917
+        self,
+        python_version: PythonVersion = PythonVersionMin,
+        use_standard_collections: bool = False,  # noqa: FBT001, FBT002
+        use_generic_container_types: bool = False,  # noqa: FBT001, FBT002
+        strict_types: Sequence[StrictTypes] | None = None,
+        use_non_positive_negative_number_constrained_types: bool = False,  # noqa: FBT001, FBT002
+        use_decimal_for_multiple_of: bool = False,  # noqa: FBT001, FBT002
+        use_union_operator: bool = False,  # noqa: FBT001, FBT002
+        use_pendulum: bool = False,  # noqa: FBT001, FBT002
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
+        treat_dot_as_module: bool = False,  # noqa: FBT001, FBT002
+        use_serialize_as_any: bool = False,  # noqa: FBT001, FBT002
+    ) -> None:
+        """Initialize type manager with datetime type mapping."""
+        super().__init__(
+            python_version,
+            use_standard_collections,
+            use_generic_container_types,
+            strict_types,
+            use_non_positive_negative_number_constrained_types,
+            use_decimal_for_multiple_of,
+            use_union_operator,
+            use_pendulum,
+            target_datetime_class,
+            treat_dot_as_module,
+            use_serialize_as_any,
+        )
+
+        datetime_map = (
+            {
+                Types.time: self.data_type.from_import(IMPORT_TIME),
+                Types.date: self.data_type.from_import(IMPORT_DATE),
+                Types.date_time: self.data_type.from_import(IMPORT_DATETIME),
+                Types.timedelta: self.data_type.from_import(IMPORT_TIMEDELTA),
+            }
+            if target_datetime_class is DatetimeClassType.Datetime
+            else {}
+        )
+
+        self.type_map: dict[Types, DataType] = {
+            **type_map_factory(self.data_type),
+            **datetime_map,
+        }

datamodel_code_generator/model/enum.py

@@ -1,38 +1,159 @@
-from typing import Any, ClassVar, Optional, Tuple
+"""Enumeration model generator.
 
-from datamodel_code_generator.imports import IMPORT_ANY, IMPORT_ENUM, Import
+Provides Enum, StrEnum, and specialized enum classes for code generation.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
+
+from datamodel_code_generator.imports import IMPORT_ANY, IMPORT_ENUM, IMPORT_INT_ENUM, IMPORT_STR_ENUM, Import
 from datamodel_code_generator.model import DataModel, DataModelFieldBase
+from datamodel_code_generator.model.base import UNDEFINED, BaseClassDataType
 from datamodel_code_generator.types import DataType, Types
 
+if TYPE_CHECKING:
+    from collections import defaultdict
+    from pathlib import Path
+
+    from datamodel_code_generator.reference import Reference
+
+
+_INT: str = "int"
+_FLOAT: str = "float"
+_BYTES: str = "bytes"
+_STR: str = "str"
+
+SUBCLASS_BASE_CLASSES: dict[Types, str] = {
+    Types.int32: _INT,
+    Types.int64: _INT,
+    Types.integer: _INT,
+    Types.float: _FLOAT,
+    Types.double: _FLOAT,
+    Types.number: _FLOAT,
+    Types.byte: _BYTES,
+    Types.string: _STR,
+}
+
 
 class Enum(DataModel):
-    TEMPLATE_FILE_PATH: ClassVar[str] = 'Enum.jinja2'
-    BASE_CLASS: ClassVar[str] = 'enum.Enum'
-    DEFAULT_IMPORTS: ClassVar[Tuple[Import, ...]] = (IMPORT_ENUM,)
+    """DataModel implementation for Python enumerations."""
+
+    TEMPLATE_FILE_PATH: ClassVar[str] = "Enum.jinja2"
+    BASE_CLASS: ClassVar[str] = "enum.Enum"
+    DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = (IMPORT_ENUM,)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        reference: Reference,
+        fields: list[DataModelFieldBase],
+        decorators: list[str] | None = None,
+        base_classes: list[Reference] | None = None,
+        custom_base_class: str | None = None,
+        custom_template_dir: Path | None = None,
+        extra_template_data: defaultdict[str, dict[str, Any]] | None = None,
+        methods: list[str] | None = None,
+        path: Path | None = None,
+        description: str | None = None,
+        type_: Types | None = None,
+        default: Any = UNDEFINED,
+        nullable: bool = False,
+        keyword_only: bool = False,
+        treat_dot_as_module: bool = False,
+    ) -> None:
+        """Initialize Enum with optional specialized base class based on type."""
+        super().__init__(
+            reference=reference,
+            fields=fields,
+            decorators=decorators,
+            base_classes=base_classes,
+            custom_base_class=custom_base_class,
+            custom_template_dir=custom_template_dir,
+            extra_template_data=extra_template_data,
+            methods=methods,
+            path=path,
+            description=description,
+            default=default,
+            nullable=nullable,
+            keyword_only=keyword_only,
+            treat_dot_as_module=treat_dot_as_module,
+        )
+        if not base_classes and type_ and (base_class := SUBCLASS_BASE_CLASSES.get(type_)):
+            self.base_classes: list[BaseClassDataType] = [
+                BaseClassDataType(type=base_class),
+                *self.base_classes,
+            ]
 
     @classmethod
     def get_data_type(cls, types: Types, **kwargs: Any) -> DataType:
+        """Get data type for enum (not implemented)."""
         raise NotImplementedError
 
-    def get_member(self, field: DataModelFieldBase) -> 'Member':
+    def get_member(self, field: DataModelFieldBase) -> Member:
+        """Create a Member instance for the given field."""
         return Member(self, field)
 
-    def find_member(self, value: Any) -> Optional['Member']:
+    def find_member(self, value: Any) -> Member | None:
+        """Find enum member matching the given value."""
         repr_value = repr(value)
-        for field in self.fields:  # pragma: no cover
-            if field.default == repr_value:
+        # Remove surrounding quotes from the string representation
+        str_value = str(value).strip("'\"")
+
+        for field in self.fields:
+            # Remove surrounding quotes from field default value
+            field_default = str(field.default or "").strip("'\"")
+
+            # Compare values after removing quotes
+            if field_default == str_value:
                 return self.get_member(field)
-        return None  # pragma: no cover
+
+            # Keep original comparison for backwards compatibility
+            if field.default == repr_value:  # pragma: no cover
+                return self.get_member(field)
+
+        return None
 
     @property
-    def imports(self) -> Tuple[Import, ...]:
+    def imports(self) -> tuple[Import, ...]:
+        """Get imports excluding Any."""
         return tuple(i for i in super().imports if i != IMPORT_ANY)
 
 
+class StrEnum(Enum):
+    """String enumeration type."""
+
+    BASE_CLASS: ClassVar[str] = "enum.StrEnum"
+    DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = (IMPORT_STR_ENUM,)
+
+
+class IntEnum(Enum):
+    """Integer enumeration type."""
+
+    BASE_CLASS: ClassVar[str] = "enum.IntEnum"
+    DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = (IMPORT_INT_ENUM,)
+
+
+SPECIALIZED_ENUM_TYPE_MATCH: dict[Types, type[Enum]] = {
+    Types.int32: IntEnum,
+    Types.int64: IntEnum,
+    Types.integer: IntEnum,
+    Types.string: StrEnum,
+}
+"""
+Map specialized enum types to their corresponding Enum subclasses.
+"""
+
+
 class Member:
+    """Represents an enum member with its parent enum and field."""
+
     def __init__(self, enum: Enum, field: DataModelFieldBase) -> None:
+        """Initialize enum member."""
         self.enum: Enum = enum
         self.field: DataModelFieldBase = field
+        self.alias: Optional[str] = None  # noqa: UP045
 
     def __repr__(self) -> str:
-        return f'{self.enum.name}.{self.field.name}'
+        """Return string representation of enum member."""
+        return f"{self.alias or self.enum.class_name}.{self.field.name}"

datamodel_code_generator/model/imports.py

@@ -0,0 +1,22 @@
+"""Import definitions for model modules.
+
+Provides pre-defined Import objects for dataclasses, TypedDict, msgspec, etc.
+"""
+
+from __future__ import annotations
+
+from datamodel_code_generator.imports import Import
+
+IMPORT_DATACLASS = Import.from_full_path("dataclasses.dataclass")
+IMPORT_FIELD = Import.from_full_path("dataclasses.field")
+IMPORT_CLASSVAR = Import.from_full_path("typing.ClassVar")
+IMPORT_TYPED_DICT = Import.from_full_path("typing.TypedDict")
+IMPORT_TYPED_DICT_BACKPORT = Import.from_full_path("typing_extensions.TypedDict")
+IMPORT_NOT_REQUIRED = Import.from_full_path("typing.NotRequired")
+IMPORT_NOT_REQUIRED_BACKPORT = Import.from_full_path("typing_extensions.NotRequired")
+IMPORT_MSGSPEC_STRUCT = Import.from_full_path("msgspec.Struct")
+IMPORT_MSGSPEC_FIELD = Import.from_full_path("msgspec.field")
+IMPORT_MSGSPEC_META = Import.from_full_path("msgspec.Meta")
+IMPORT_MSGSPEC_CONVERT = Import.from_full_path("msgspec.convert")
+IMPORT_MSGSPEC_UNSET = Import.from_full_path("msgspec.UNSET")
+IMPORT_MSGSPEC_UNSETTYPE = Import.from_full_path("msgspec.UnsetType")