Support datetime types in dataclass fields

applying `--output-datetime-class` from #2100 to dataclass to map date, time and date time to the python `datetime` objects instead of strings.
koxudaxi · Oct 15, 2024 · 4d67a81 · 4d67a81
1 parent b2899e1
commit 4d67a81
Show file tree

Hide file tree

Showing 12 changed files with 116 additions and 25 deletions.
diff --git a/README.md b/README.md
@@ -452,6 +452,10 @@ Model customization:
   --keep-model-order    Keep generated models'' order
   --keyword-only        Defined models as keyword only (for example
                         dataclass(kw_only=True)).
+  --output-datetime-class {datetime,AwareDatetime,NaiveDatetime}
+                        Choose Datetime class between AwareDatetime, NaiveDatetime or
+                        datetime. Each output model has its default mapping, and only
+                        pydantic and dataclass support this override"
   --reuse-model         Reuse models on the field when a module has the model with the same
                         content
   --target-python-version {3.6,3.7,3.8,3.9,3.10,3.11,3.12}
@@ -464,8 +468,6 @@ Model customization:
   --use-schema-description
                         Use schema description to populate class docstring
   --use-title-as-name   use titles as class names of models
-
-  ----output-datetime-class Choose Datetime class between AwareDatetime, NaiveDatetime or datetime, default: "datetime"
 
 Template customization:
   --aliases ALIASES     Alias mapping file

diff --git a/datamodel_code_generator/__init__.py b/datamodel_code_generator/__init__.py
@@ -301,7 +301,7 @@ def generate(
     treat_dots_as_module: bool = False,
     use_exact_imports: bool = False,
     union_mode: Optional[UnionMode] = None,
-    output_datetime_class: DataModelType = DatetimeClassType.Datetime,
+    output_datetime_class: Optional[DatetimeClassType] = None,
     keyword_only: bool = False,
 ) -> None:
     remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()

diff --git a/datamodel_code_generator/__main__.py b/datamodel_code_generator/__main__.py
@@ -193,6 +193,22 @@ def validate_keyword_only(cls, values: Dict[str, Any]) -> Dict[str, Any]:
             )
         return values
 
+    @model_validator(mode='after')
+    def validate_output_datetime_class(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        datetime_class_type: Optional[DatetimeClassType] = values.get(
+            'output_datetime_class'
+        )
+        if (
+            datetime_class_type
+            and datetime_class_type is not DatetimeClassType.Datetime
+            and values.get('output_model_type') == DataModelType.DataclassesDataclass
+        ):
+            raise Error(
+                '`--output-datetime-class` only allows "datetime" for '
+                f'`--output-model-type` {DataModelType.DataclassesDataclass.value}'
+            )
+        return values
+
     # Pydantic 1.5.1 doesn't support each_item=True correctly
     @field_validator('http_headers', mode='before')
     def validate_http_headers(cls, value: Any) -> Optional[List[Tuple[str, str]]]:
@@ -323,7 +339,7 @@ def validate_root(cls, values: Any) -> Any:
     treat_dot_as_module: bool = False
     use_exact_imports: bool = False
     union_mode: Optional[UnionMode] = None
-    output_datetime_class: DatetimeClassType = DatetimeClassType.Datetime
+    output_datetime_class: Optional[DatetimeClassType] = None
     keyword_only: bool = False
 
     def merge_args(self, args: Namespace) -> None:

diff --git a/datamodel_code_generator/arguments.py b/datamodel_code_generator/arguments.py
@@ -200,8 +200,10 @@ def start_section(self, heading: Optional[str]) -> None:
 )
 model_options.add_argument(
     '--output-datetime-class',
-    help='Choose Datetime class between AwareDatetime, NaiveDatetime or datetime, default: "datetime"',
+    help='Choose Datetime class between AwareDatetime, NaiveDatetime or datetime. '
+    'Each output model has its default mapping (for example pydantic: datetime, dataclass: str, ...)',
     choices=[i.value for i in DatetimeClassType],
+    default=None,
 )
 
 # ======================================================================================

diff --git a/datamodel_code_generator/model/__init__.py b/datamodel_code_generator/model/__init__.py
@@ -48,7 +48,7 @@ def get_data_model_types(
             data_model=dataclass.DataClass,
             root_model=rootmodel.RootModel,
             field_model=dataclass.DataModelField,
-            data_type_manager=DataTypeManager,
+            data_type_manager=dataclass.DataTypeManager,
             dump_resolve_reference_action=None,
         )
     elif data_model_type == DataModelType.TypingTypedDict:

diff --git a/datamodel_code_generator/model/dataclass.py b/datamodel_code_generator/model/dataclass.py
@@ -1,13 +1,32 @@
 from pathlib import Path
-from typing import Any, ClassVar, DefaultDict, Dict, List, Optional, Set, Tuple
-
-from datamodel_code_generator.imports import Import
+from typing import (
+    Any,
+    ClassVar,
+    DefaultDict,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+)
+
+from datamodel_code_generator import DatetimeClassType, PythonVersion
+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
+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 chain_as_tuple
+from datamodel_code_generator.types import DataType, StrictTypes, Types, chain_as_tuple
 
 
 def _has_field_assignment(field: DataModelFieldBase) -> bool:
@@ -120,3 +139,43 @@ def __str__(self) -> str:
             f'{k}={v if k == "default_factory" else repr(v)}' for k, v in data.items()
         ]
         return f'field({", ".join(kwargs)})'
+
+
+class DataTypeManager(_DataTypeManager):
+    def __init__(
+        self,
+        python_version: PythonVersion = PythonVersion.PY_38,
+        use_standard_collections: bool = False,
+        use_generic_container_types: bool = False,
+        strict_types: Optional[Sequence[StrictTypes]] = None,
+        use_non_positive_negative_number_constrained_types: bool = False,
+        use_union_operator: bool = False,
+        use_pendulum: bool = False,
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
+    ):
+        super().__init__(
+            python_version,
+            use_standard_collections,
+            use_generic_container_types,
+            strict_types,
+            use_non_positive_negative_number_constrained_types,
+            use_union_operator,
+            use_pendulum,
+            target_datetime_class,
+        )
+
+        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,
+        }
diff --git a/datamodel_code_generator/model/pydantic/types.py b/datamodel_code_generator/model/pydantic/types.py
@@ -163,7 +163,7 @@ def __init__(
         use_non_positive_negative_number_constrained_types: bool = False,
         use_union_operator: bool = False,
         use_pendulum: bool = False,
-        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
+        target_datetime_class: Optional[DatetimeClassType] = None,
     ):
         super().__init__(
             python_version,

diff --git a/datamodel_code_generator/model/pydantic_v2/types.py b/datamodel_code_generator/model/pydantic_v2/types.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import ClassVar, Dict, Sequence, Type
+from typing import ClassVar, Dict, Optional, Sequence, Type
 
 from datamodel_code_generator.format import DatetimeClassType
 from datamodel_code_generator.model.pydantic import DataTypeManager as _DataTypeManager
@@ -20,7 +20,7 @@ def type_map_factory(
         data_type: Type[DataType],
         strict_types: Sequence[StrictTypes],
         pattern_key: str,
-        target_datetime_class: DatetimeClassType,
+        target_datetime_class: Optional[DatetimeClassType] = None,
     ) -> Dict[Types, DataType]:
         result = {
             **super().type_map_factory(

diff --git a/datamodel_code_generator/types.py b/datamodel_code_generator/types.py
@@ -575,7 +575,7 @@ def __init__(
         use_non_positive_negative_number_constrained_types: bool = False,
         use_union_operator: bool = False,
         use_pendulum: bool = False,
-        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
+        target_datetime_class: Optional[DatetimeClassType] = None,
     ) -> None:
         self.python_version = python_version
         self.use_standard_collections: bool = use_standard_collections

diff --git a/docs/index.md b/docs/index.md
@@ -446,6 +446,10 @@ Model customization:
   --keep-model-order    Keep generated models'' order
   --keyword-only        Defined models as keyword only (for example
                         dataclass(kw_only=True)).
+  --output-datetime-class {datetime,AwareDatetime,NaiveDatetime}
+                        Choose Datetime class between AwareDatetime, NaiveDatetime or
+                        datetime. Each output model has its default mapping, and only
+                        pydantic and dataclass support this override"
   --reuse-model         Reuse models on the field when a module has the model with the same
                         content
   --target-python-version {3.6,3.7,3.8,3.9,3.10,3.11,3.12}

diff --git a/tests/data/expected/main/openapi/datetime_dataclass.py b/tests/data/expected/main/openapi/datetime_dataclass.py
@@ -0,0 +1,13 @@
+# generated by datamodel-codegen:
+#   filename:  datetime.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass
+class InventoryItem:
+    releaseDate: datetime
diff --git a/tests/main/openapi/test_main_openapi.py b/tests/main/openapi/test_main_openapi.py
@@ -1075,19 +1075,14 @@ def test_main_original_field_name_delimiter_without_snake_case_field(capsys) ->
 
 @freeze_time('2019-07-26')
 @pytest.mark.parametrize(
-    'output_model,expected_output',
+    'output_model,expected_output,date_type',
     [
-        (
-            'pydantic.BaseModel',
-            'datetime.py',
-        ),
-        (
-            'pydantic_v2.BaseModel',
-            'datetime_pydantic_v2.py',
-        ),
+        ('pydantic.BaseModel', 'datetime.py', 'AwareDatetime'),
+        ('pydantic_v2.BaseModel', 'datetime_pydantic_v2.py', 'AwareDatetime'),
+        ('dataclasses.dataclass', 'datetime_dataclass.py', 'datetime'),
     ],
 )
-def test_main_openapi_aware_datetime(output_model, expected_output):
+def test_main_openapi_aware_datetime(output_model, expected_output, date_type):
     with TemporaryDirectory() as output_dir:
         output_file: Path = Path(output_dir) / 'output.py'
         return_code: Exit = main(
@@ -1099,7 +1094,7 @@ def test_main_openapi_aware_datetime(output_model, expected_output):
                 '--input-file-type',
                 'openapi',
                 '--output-datetime-class',
-                'AwareDatetime',
+                date_type,
                 '--output-model',
                 output_model,
             ]