From 4d67a81030f9058e9e292f46ca59bbd396d07ac2 Mon Sep 17 00:00:00 2001 From: Anis Da Silva Campos Date: Sat, 12 Oct 2024 18:52:31 +0200 Subject: [PATCH] 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. --- README.md | 6 +- datamodel_code_generator/__init__.py | 2 +- datamodel_code_generator/__main__.py | 18 ++++- datamodel_code_generator/arguments.py | 4 +- datamodel_code_generator/model/__init__.py | 2 +- datamodel_code_generator/model/dataclass.py | 67 +++++++++++++++++-- .../model/pydantic/types.py | 2 +- .../model/pydantic_v2/types.py | 4 +- datamodel_code_generator/types.py | 2 +- docs/index.md | 4 ++ .../main/openapi/datetime_dataclass.py | 13 ++++ tests/main/openapi/test_main_openapi.py | 17 ++--- 12 files changed, 116 insertions(+), 25 deletions(-) create mode 100644 tests/data/expected/main/openapi/datetime_dataclass.py diff --git a/README.md b/README.md index 06530557..740c2eb1 100644 --- 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 index 2f4c0033..da52789e 100644 --- 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 index de59993a..258332e4 100644 --- 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 index 6706b44b..3b4ff421 100644 --- 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 index 19dcd1e0..79e3d1fb 100644 --- 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 index b4f17b9d..fd492d61 100644 --- 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 index 7938ef74..e9fc37f6 100644 --- 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 index 9d901348..d0327fcc 100644 --- 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 index 498bbc22..6df57e7d 100644 --- 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 index b4aa656d..485e5c95 100644 --- 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 new file mode 100644 index 00000000..e6b77879 --- /dev/null +++ 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 index 0a42502b..08761bba 100644 --- 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, ]