with pydantic v2, use pydantic.AwareDatetime instead of datetime (#1735)

* with pydantic v2, use pydantic.AwareDatetime instead of datetime Openapi spec says that type string with format "date-time" must have a timezone. Therefore, we need pydantic.AwareDatetime here. * Add unittest --------- Co-authored-by: Koudai Aono <[email protected]>
koxudaxi · Nov 24, 2023 · a31148d · a31148d
1 parent a36ce94
commit a31148d
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 2 deletions.
diff --git a/datamodel_code_generator/model/pydantic_v2/imports.py b/datamodel_code_generator/model/pydantic_v2/imports.py
@@ -1,3 +1,4 @@
 from datamodel_code_generator.imports import Import
 
 IMPORT_CONFIG_DICT = Import.from_full_path('pydantic.ConfigDict')
+IMPORT_AWARE_DATETIME = Import.from_full_path('pydantic.AwareDatetime')
diff --git a/datamodel_code_generator/model/pydantic_v2/types.py b/datamodel_code_generator/model/pydantic_v2/types.py
@@ -4,6 +4,7 @@
 
 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.types import DataType, StrictTypes, Types
 
 
@@ -27,4 +28,5 @@ def type_map_factory(
                     **({'strict': True} if StrictTypes.str in strict_types else {}),
                 },
             ),
+            Types.date_time: data_type.from_import(IMPORT_AWARE_DATETIME),
         }
diff --git a/tests/data/expected/main/main_openapi_datetime_pydantic_v2/output.py b/tests/data/expected/main/main_openapi_datetime_pydantic_v2/output.py
@@ -0,0 +1,11 @@
+# generated by datamodel-codegen:
+#   filename:  datetime.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import AwareDatetime, BaseModel, Field
+
+
+class InventoryItem(BaseModel):
+    releaseDate: AwareDatetime = Field(..., examples=['2016-08-29T09:12:33.001Z'])
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -2067,7 +2067,20 @@ def test_main_json_capitalise_enum_members_without_enum():
 
 
 @freeze_time('2019-07-26')
-def test_main_openapi_datetime():
+@pytest.mark.parametrize(
+    'output_model,expected_output',
+    [
+        (
+            'pydantic.BaseModel',
+            'main_openapi_datetime',
+        ),
+        (
+            'pydantic_v2.BaseModel',
+            'main_openapi_datetime_pydantic_v2',
+        ),
+    ],
+)
+def test_main_openapi_datetime(output_model, expected_output):
     with TemporaryDirectory() as output_dir:
         output_file: Path = Path(output_dir) / 'output.py'
         return_code: Exit = main(
@@ -2078,12 +2091,14 @@ def test_main_openapi_datetime():
                 str(output_file),
                 '--input-file-type',
                 'openapi',
+                '--output-model',
+                output_model,
             ]
         )
         assert return_code == Exit.OK
         assert (
             output_file.read_text()
-            == (EXPECTED_MAIN_PATH / 'main_openapi_datetime' / 'output.py').read_text()
+            == (EXPECTED_MAIN_PATH / expected_output / 'output.py').read_text()
         )