New CLI option --output-datetime-class #1996 (#2100)

* [IMP] new CLI option --output-datetime-class to choose between AwareDateTime, NaiveDateTime, generic datetime #1996

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [IMP]  tests and README

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

README.md

@@ -462,6 +462,8 @@ 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

datamodel_code_generator/__init__.py

@@ -30,7 +30,7 @@
 import yaml
 
 import datamodel_code_generator.pydantic_patch  # noqa: F401
-from datamodel_code_generator.format import PythonVersion
+from datamodel_code_generator.format import DatetimeClassType, PythonVersion
 from datamodel_code_generator.model.pydantic_v2 import UnionMode
 from datamodel_code_generator.parser import DefaultPutDict, LiteralType
 from datamodel_code_generator.parser.base import Parser
@@ -301,6 +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,
 ) -> None:
     remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()
     if isinstance(input_, str):
@@ -395,9 +396,12 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
             raise Error('union_mode is only supported for pydantic_v2.BaseModel')
     else:
         default_field_extras = None
+
     from datamodel_code_generator.model import get_data_model_types
 
-    data_model_types = get_data_model_types(output_model_type, target_python_version)
+    data_model_types = get_data_model_types(
+        output_model_type, target_python_version, output_datetime_class
+    )
     parser = parser_class(
         source=input_text or input_,
         data_model_type=data_model_types.data_model,
@@ -471,6 +475,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
         treat_dots_as_module=treat_dots_as_module,
         use_exact_imports=use_exact_imports,
         default_field_extras=default_field_extras,
+        target_datetime_class=output_datetime_class,
         **kwargs,
     )
 

datamodel_code_generator/__main__.py

@@ -51,6 +51,7 @@
 )
 from datamodel_code_generator.arguments import DEFAULT_ENCODING, arg_parser, namespace
 from datamodel_code_generator.format import (
+    DatetimeClassType,
     PythonVersion,
     black_find_project_root,
     is_supported_in_black,
@@ -313,6 +314,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
 
     def merge_args(self, args: Namespace) -> None:
         set_args = {
@@ -512,6 +514,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit:
             treat_dots_as_module=config.treat_dot_as_module,
             use_exact_imports=config.use_exact_imports,
             union_mode=config.union_mode,
+            output_datetime_class=config.output_datetime_class,
         )
         return Exit.OK
     except InvalidClassNameError as e:

datamodel_code_generator/arguments.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from datamodel_code_generator import DataModelType, InputFileType, OpenAPIScope
-from datamodel_code_generator.format import PythonVersion
+from datamodel_code_generator.format import DatetimeClassType, PythonVersion
 from datamodel_code_generator.model.pydantic_v2 import UnionMode
 from datamodel_code_generator.parser import LiteralType
 from datamodel_code_generator.types import StrictTypes
@@ -192,6 +192,11 @@ def start_section(self, heading: Optional[str]) -> None:
     action='store_true',
     default=False,
 )
+model_options.add_argument(
+    '--output-datetime-class',
+    help='Choose Datetime class between AwareDatetime, NaiveDatetime or datetime, default: "datetime"',
+    choices=[i.value for i in DatetimeClassType],
+)
 
 # ======================================================================================
 # Typing options for generated models

datamodel_code_generator/format.py

@@ -17,6 +17,12 @@
     black.mode = None
 
 
+class DatetimeClassType(Enum):
+    Datetime = 'datetime'
+    Awaredatetime = 'AwareDatetime'
+    Naivedatetime = 'NaiveDatetime'
+
+
 class PythonVersion(Enum):
     PY_36 = '3.6'
     PY_37 = '3.7'

datamodel_code_generator/model/__init__.py

@@ -6,7 +6,7 @@
 from .base import ConstraintsBase, DataModel, DataModelFieldBase
 
 if TYPE_CHECKING:
-    from .. import DataModelType, PythonVersion
+    from .. import DataModelType, DatetimeClassType, PythonVersion
 
 
 class DataModelSet(NamedTuple):
@@ -19,7 +19,9 @@ class DataModelSet(NamedTuple):
 
 
 def get_data_model_types(
-    data_model_type: DataModelType, target_python_version: PythonVersion
+    data_model_type: DataModelType,
+    target_python_version: PythonVersion,
+    target_datetime_class: DatetimeClassType,
 ) -> DataModelSet:
     from .. import DataModelType
     from . import dataclass, msgspec, pydantic, pydantic_v2, rootmodel, typed_dict

datamodel_code_generator/model/pydantic/types.py

@@ -3,7 +3,7 @@
 from decimal import Decimal
 from typing import Any, ClassVar, Dict, Optional, Sequence, Set, Type
 
-from datamodel_code_generator.format import PythonVersion
+from datamodel_code_generator.format import DatetimeClassType, PythonVersion
 from datamodel_code_generator.imports import (
     IMPORT_ANY,
     IMPORT_DATE,
@@ -59,6 +59,7 @@ def type_map_factory(
     strict_types: Sequence[StrictTypes],
     pattern_key: str,
     use_pendulum: bool,
+    target_datetime_class: DatetimeClassType,
 ) -> Dict[Types, DataType]:
     data_type_int = data_type(type='int')
     data_type_float = data_type(type='float')
@@ -162,6 +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,
     ):
         super().__init__(
             python_version,
@@ -171,12 +173,14 @@ def __init__(
             use_non_positive_negative_number_constrained_types,
             use_union_operator,
             use_pendulum,
+            target_datetime_class,
         )
 
         self.type_map: Dict[Types, DataType] = self.type_map_factory(
             self.data_type,
             strict_types=self.strict_types,
             pattern_key=self.PATTERN_KEY,
+            target_datetime_class=target_datetime_class,
         )
         self.strict_type_map: Dict[StrictTypes, DataType] = strict_type_map_factory(
             self.data_type,
@@ -200,8 +204,15 @@ def type_map_factory(
         data_type: Type[DataType],
         strict_types: Sequence[StrictTypes],
         pattern_key: str,
+        target_datetime_class: DatetimeClassType,
     ) -> Dict[Types, DataType]:
-        return type_map_factory(data_type, strict_types, pattern_key, self.use_pendulum)
+        return type_map_factory(
+            data_type,
+            strict_types,
+            pattern_key,
+            self.use_pendulum,
+            self.target_datetime_class,
+        )
 
     def transform_kwargs(
         self, kwargs: Dict[str, Any], filter_: Set[str]

datamodel_code_generator/model/pydantic_v2/imports.py

@@ -2,3 +2,4 @@
 
 IMPORT_CONFIG_DICT = Import.from_full_path('pydantic.ConfigDict')
 IMPORT_AWARE_DATETIME = Import.from_full_path('pydantic.AwareDatetime')
+IMPORT_NAIVE_DATETIME = Import.from_full_path('pydantic.NaiveDatetime')

datamodel_code_generator/model/pydantic_v2/types.py

@@ -2,9 +2,13 @@
 
 from typing import ClassVar, Dict, Sequence, Type
 
+from datamodel_code_generator.format import DatetimeClassType
 from datamodel_code_generator.model.pydantic import DataTypeManager as _DataTypeManager
 from datamodel_code_generator.model.pydantic.imports import IMPORT_CONSTR
-from datamodel_code_generator.model.pydantic_v2.imports import IMPORT_AWARE_DATETIME
+from datamodel_code_generator.model.pydantic_v2.imports import (
+    IMPORT_AWARE_DATETIME,
+    IMPORT_NAIVE_DATETIME,
+)
 from datamodel_code_generator.types import DataType, StrictTypes, Types
 
 
@@ -16,9 +20,12 @@ def type_map_factory(
         data_type: Type[DataType],
         strict_types: Sequence[StrictTypes],
         pattern_key: str,
+        target_datetime_class: DatetimeClassType,
     ) -> Dict[Types, DataType]:
-        return {
-            **super().type_map_factory(data_type, strict_types, pattern_key),
+        result = {
+            **super().type_map_factory(
+                data_type, strict_types, pattern_key, target_datetime_class
+            ),
             Types.hostname: self.data_type.from_import(
                 IMPORT_CONSTR,
                 strict=StrictTypes.str in strict_types,
@@ -28,5 +35,9 @@ def type_map_factory(
                     **({'strict': True} if StrictTypes.str in strict_types else {}),
                 },
             ),
-            Types.date_time: data_type.from_import(IMPORT_AWARE_DATETIME),
         }
+        if target_datetime_class == DatetimeClassType.Awaredatetime:
+            result[Types.date_time] = data_type.from_import(IMPORT_AWARE_DATETIME)
+        if target_datetime_class == DatetimeClassType.Naivedatetime:
+            result[Types.date_time] = data_type.from_import(IMPORT_NAIVE_DATETIME)
+        return result

datamodel_code_generator/model/types.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional, Sequence, Type
 
-from datamodel_code_generator import PythonVersion
+from datamodel_code_generator import DatetimeClassType, PythonVersion
 from datamodel_code_generator.imports import (
     IMPORT_ANY,
     IMPORT_DECIMAL,
@@ -10,9 +10,7 @@
 from datamodel_code_generator.types import DataTypeManager as _DataTypeManager
 
 
-def type_map_factory(
-    data_type: Type[DataType],
-) -> Dict[Types, DataType]:
+def type_map_factory(data_type: Type[DataType]) -> Dict[Types, DataType]:
     data_type_int = data_type(type='int')
     data_type_float = data_type(type='float')
     data_type_str = data_type(type='str')
@@ -64,6 +62,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,
     ):
         super().__init__(
             python_version,
@@ -73,11 +72,10 @@ def __init__(
             use_non_positive_negative_number_constrained_types,
             use_union_operator,
             use_pendulum,
+            target_datetime_class,
         )
 
-        self.type_map: Dict[Types, DataType] = type_map_factory(
-            self.data_type,
-        )
+        self.type_map: Dict[Types, DataType] = type_map_factory(self.data_type)
 
     def get_data_type(
         self,

datamodel_code_generator/parser/base.py

@@ -26,7 +26,11 @@
 
 from pydantic import BaseModel
 
-from datamodel_code_generator.format import CodeFormatter, PythonVersion
+from datamodel_code_generator.format import (
+    CodeFormatter,
+    DatetimeClassType,
+    PythonVersion,
+)
 from datamodel_code_generator.imports import (
     IMPORT_ANNOTATIONS,
     IMPORT_LITERAL,
@@ -404,6 +408,7 @@ def __init__(
         treat_dots_as_module: bool = False,
         use_exact_imports: bool = False,
         default_field_extras: Optional[Dict[str, Any]] = None,
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
     ) -> None:
         self.data_type_manager: DataTypeManager = data_type_manager_type(
             python_version=target_python_version,
@@ -412,6 +417,7 @@ def __init__(
             strict_types=strict_types,
             use_union_operator=use_union_operator,
             use_pendulum=use_pendulum,
+            target_datetime_class=target_datetime_class,
         )
         self.data_model_type: Type[DataModel] = data_model_type
         self.data_model_root_type: Type[DataModel] = data_model_root_type

datamodel_code_generator/parser/graphql.py

@@ -46,6 +46,7 @@
         "Please run `$pip install 'datamodel-code-generator[graphql]`' to generate data-model from a GraphQL schema."
     )
 
+from datamodel_code_generator.format import DatetimeClassType
 
 graphql_resolver = graphql.type.introspection.TypeResolvers()
 
@@ -157,6 +158,7 @@ def __init__(
         treat_dots_as_module: bool = False,
         use_exact_imports: bool = False,
         default_field_extras: Optional[Dict[str, Any]] = None,
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
     ) -> None:
         super().__init__(
             source=source,
@@ -227,6 +229,7 @@ def __init__(
             treat_dots_as_module=treat_dots_as_module,
             use_exact_imports=use_exact_imports,
             default_field_extras=default_field_extras,
+            target_datetime_class=target_datetime_class,
         )
 
         self.data_model_scalar_type = data_model_scalar_type

datamodel_code_generator/parser/jsonschema.py

@@ -71,6 +71,8 @@
 if PYDANTIC_V2:
     from pydantic import ConfigDict
 
+from datamodel_code_generator.format import DatetimeClassType
+
 
 def get_model_by_path(
     schema: Union[Dict[str, Any], List[Any]], keys: Union[List[str], List[int]]
@@ -444,6 +446,7 @@ def __init__(
         treat_dots_as_module: bool = False,
         use_exact_imports: bool = False,
         default_field_extras: Optional[Dict[str, Any]] = None,
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
     ) -> None:
         super().__init__(
             source=source,
@@ -514,6 +517,7 @@ def __init__(
             treat_dots_as_module=treat_dots_as_module,
             use_exact_imports=use_exact_imports,
             default_field_extras=default_field_extras,
+            target_datetime_class=target_datetime_class,
         )
 
         self.remote_object_cache: DefaultPutDict[str, Dict[str, Any]] = DefaultPutDict()

datamodel_code_generator/parser/openapi.py

@@ -36,6 +36,7 @@
     load_yaml,
     snooper_to_methods,
 )
+from datamodel_code_generator.format import DatetimeClassType
 from datamodel_code_generator.model import DataModel, DataModelFieldBase
 from datamodel_code_generator.model import pydantic as pydantic_model
 from datamodel_code_generator.parser.base import get_special_path
@@ -225,6 +226,7 @@ def __init__(
         treat_dots_as_module: bool = False,
         use_exact_imports: bool = False,
         default_field_extras: Optional[Dict[str, Any]] = None,
+        target_datetime_class: DatetimeClassType = DatetimeClassType.Datetime,
     ):
         super().__init__(
             source=source,
@@ -295,6 +297,7 @@ def __init__(
             treat_dots_as_module=treat_dots_as_module,
             use_exact_imports=use_exact_imports,
             default_field_extras=default_field_extras,
+            target_datetime_class=target_datetime_class,
         )
         self.open_api_scopes: List[OpenAPIScope] = openapi_scopes or [
             OpenAPIScope.Schemas

datamodel_code_generator/types.py

@@ -27,7 +27,7 @@
 from packaging import version
 from pydantic import StrictBool, StrictInt, StrictStr, create_model
 
-from datamodel_code_generator.format import PythonVersion
+from datamodel_code_generator.format import DatetimeClassType, PythonVersion
 from datamodel_code_generator.imports import (
     IMPORT_ABC_MAPPING,
     IMPORT_ABC_SEQUENCE,
@@ -575,6 +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,
     ) -> None:
         self.python_version = python_version
         self.use_standard_collections: bool = use_standard_collections
@@ -585,6 +586,7 @@ def __init__(
         )
         self.use_union_operator: bool = use_union_operator
         self.use_pendulum: bool = use_pendulum
+        self.target_datetime_class: DatetimeClassType = target_datetime_class
 
         if (
             use_generic_container_types and python_version == PythonVersion.PY_36