Skip to content

Commit

Permalink
New CLI option --output-datetime-class #1996 (#2100)
Browse files Browse the repository at this point in the history
* [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>
  • Loading branch information
archetipo and pre-commit-ci[bot] authored Oct 12, 2024
1 parent f11c371 commit 2df133c
Show file tree
Hide file tree
Showing 21 changed files with 619 additions and 367 deletions.
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 7 additions & 2 deletions datamodel_code_generator/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
)

Expand Down
3 changes: 3 additions & 0 deletions datamodel_code_generator/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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:
Expand Down
7 changes: 6 additions & 1 deletion datamodel_code_generator/arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions datamodel_code_generator/format.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
6 changes: 4 additions & 2 deletions datamodel_code_generator/model/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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
Expand Down
15 changes: 13 additions & 2 deletions datamodel_code_generator/model/pydantic/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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')
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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]
Expand Down
1 change: 1 addition & 0 deletions datamodel_code_generator/model/pydantic_v2/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
19 changes: 15 additions & 4 deletions datamodel_code_generator/model/pydantic_v2/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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,
Expand All @@ -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
12 changes: 5 additions & 7 deletions datamodel_code_generator/model/types.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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')
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion datamodel_code_generator/parser/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions datamodel_code_generator/parser/graphql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions datamodel_code_generator/parser/jsonschema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down
3 changes: 3 additions & 0 deletions datamodel_code_generator/parser/openapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion datamodel_code_generator/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Loading

0 comments on commit 2df133c

Please sign in to comment.