feat: support string format: duration for pydantic/msgspec (#2084)

datamodel_code_generator/imports.py

@@ -117,9 +117,11 @@ def remove_referenced_imports(self, reference_path: str) -> None:
 IMPORT_DECIMAL = Import.from_full_path('decimal.Decimal')
 IMPORT_DATE = Import.from_full_path('datetime.date')
 IMPORT_DATETIME = Import.from_full_path('datetime.datetime')
+IMPORT_TIMEDELTA = Import.from_full_path('datetime.timedelta')
 IMPORT_PATH = Import.from_full_path('pathlib.Path')
 IMPORT_TIME = Import.from_full_path('datetime.time')
 IMPORT_UUID = Import.from_full_path('uuid.UUID')
 IMPORT_PENDULUM_DATE = Import.from_full_path('pendulum.Date')
 IMPORT_PENDULUM_DATETIME = Import.from_full_path('pendulum.DateTime')
+IMPORT_PENDULUM_DURATION = Import.from_full_path('pendulum.Duration')
 IMPORT_PENDULUM_TIME = Import.from_full_path('pendulum.Time')

datamodel_code_generator/model/pydantic/types.py

@@ -12,8 +12,10 @@
     IMPORT_PATH,
     IMPORT_PENDULUM_DATE,
     IMPORT_PENDULUM_DATETIME,
+    IMPORT_PENDULUM_DURATION,
     IMPORT_PENDULUM_TIME,
     IMPORT_TIME,
+    IMPORT_TIMEDELTA,
     IMPORT_UUID,
 )
 from datamodel_code_generator.model.pydantic.imports import (
@@ -75,6 +77,7 @@ def type_map_factory(
         Types.binary: data_type(type='bytes'),
         Types.date: data_type.from_import(IMPORT_DATE),
         Types.date_time: data_type.from_import(IMPORT_DATETIME),
+        Types.timedelta: data_type.from_import(IMPORT_TIMEDELTA),
         Types.path: data_type.from_import(IMPORT_PATH),
         Types.password: data_type.from_import(IMPORT_SECRET_STR),
         Types.email: data_type.from_import(IMPORT_EMAIL_STR),
@@ -108,6 +111,7 @@ def type_map_factory(
         result[Types.date] = data_type.from_import(IMPORT_PENDULUM_DATE)
         result[Types.date_time] = data_type.from_import(IMPORT_PENDULUM_DATETIME)
         result[Types.time] = data_type.from_import(IMPORT_PENDULUM_TIME)
+        result[Types.timedelta] = data_type.from_import(IMPORT_PENDULUM_DURATION)
 
     return result
 

datamodel_code_generator/model/types.py

@@ -1,7 +1,11 @@
 from typing import Any, Dict, Optional, Sequence, Type
 
 from datamodel_code_generator import PythonVersion
-from datamodel_code_generator.imports import IMPORT_ANY, IMPORT_DECIMAL
+from datamodel_code_generator.imports import (
+    IMPORT_ANY,
+    IMPORT_DECIMAL,
+    IMPORT_TIMEDELTA,
+)
 from datamodel_code_generator.types import DataType, StrictTypes, Types
 from datamodel_code_generator.types import DataTypeManager as _DataTypeManager
 
@@ -27,6 +31,7 @@ def type_map_factory(
         Types.binary: data_type(type='bytes'),
         Types.date: data_type_str,
         Types.date_time: data_type_str,
+        Types.timedelta: data_type.from_import(IMPORT_TIMEDELTA),
         Types.password: data_type_str,
         Types.email: data_type_str,
         Types.uuid: data_type_str,

datamodel_code_generator/parser/jsonschema.py

@@ -116,6 +116,7 @@ def get_model_by_path(
         'binary': Types.binary,
         'date': Types.date,
         'date-time': Types.date_time,
+        'duration': Types.timedelta,
         'time': Types.time,
         'password': Types.password,
         'path': Types.path,

datamodel_code_generator/types.py

@@ -542,6 +542,7 @@ class Types(Enum):
     binary = auto()
     date = auto()
     date_time = auto()
+    timedelta = auto()
     password = auto()
     path = auto()
     email = auto()

tests/data/expected/main/jsonschema/duration_msgspec.py

@@ -0,0 +1,16 @@
+# generated by datamodel-codegen:
+#   filename:  duration.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Any, Optional
+
+from msgspec import Struct
+
+Model = Any
+
+
+class Test(Struct):
+    s_duration: Optional[timedelta] = None

tests/data/expected/main/jsonschema/duration_pydantic_v2.py

@@ -0,0 +1,18 @@
+# generated by datamodel-codegen:
+#   filename:  duration.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class Model(RootModel[Any]):
+    root: Any
+
+
+class Test(BaseModel):
+    s_duration: Optional[timedelta] = Field(None, examples=['PT2H33M3S'])

tests/data/jsonschema/duration.json

@@ -0,0 +1,15 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "definitions": {
+    "Test": {
+      "type": "object",
+      "properties": {
+        "s_duration": {
+          "type": "string",
+          "format": "duration",
+          "example": "PT2H33M3S"
+        }
+      }
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -3432,3 +3432,39 @@ def test_main_imports_correct():
                 path.relative_to(main_modular_dir)
             ).read_text()
             assert result == path.read_text()
+
+
+@pytest.mark.parametrize(
+    'output_model,expected_output',
+    [
+        (
+            'pydantic_v2.BaseModel',
+            'duration_pydantic_v2.py',
+        ),
+        (
+            'msgspec.Struct',
+            'duration_msgspec.py',
+        ),
+    ],
+)
+@freeze_time('2019-07-26')
+def test_main_jsonschema_duration(output_model, expected_output):
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(JSON_SCHEMA_DATA_PATH / 'duration.json'),
+                '--output',
+                str(output_file),
+                '--output-model-type',
+                output_model,
+                '--target-python',
+                '3.8',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (EXPECTED_JSON_SCHEMA_PATH / expected_output).read_text()
+        )

tests/parser/test_jsonschema.py

@@ -411,6 +411,8 @@ def test_parse_nested_array():
         ('string', 'date', 'Date', 'pendulum', 'Date', True),
         ('string', 'date-time', 'datetime', 'datetime', 'datetime', False),
         ('string', 'date-time', 'DateTime', 'pendulum', 'DateTime', True),
+        ('string', 'duration', 'timedelta', 'datetime', 'timedelta', False),
+        ('string', 'duration', 'Duration', 'pendulum', 'Duration', True),
         ('string', 'path', 'Path', 'pathlib', 'Path', False),
         ('string', 'password', 'SecretStr', 'pydantic', 'SecretStr', False),
         ('string', 'email', 'EmailStr', 'pydantic', 'EmailStr', False),