feat: fully support pydantic v2 type handling

Author: tractorcow

src/openapi_python_generator/language_converters/python/common.py

@@ -1,9 +1,11 @@
 import keyword
 import re
 from typing import Optional
+from openapi_python_generator.common import PydanticVersion
 
 
 _use_orjson: bool = False
+_pydantic_version: PydanticVersion = PydanticVersion.V2
 _custom_template_path: str = None
 _symbol_ascii_strip_re = re.compile(r"[^A-Za-z0-9_]")
 
@@ -16,6 +18,13 @@ def set_use_orjson(value: bool) -> None:
     global _use_orjson
     _use_orjson = value
 
+def set_pydantic_version(value: PydanticVersion) -> None:
+    """
+    Set the value of the global variable
+    :param value: value of the variable
+    """
+    global _pydantic_version
+    _pydantic_version = value
 
 def get_use_orjson() -> bool:
     """
@@ -25,6 +34,13 @@ def get_use_orjson() -> bool:
     global _use_orjson
     return _use_orjson
 
+def get_pydantic_version() -> PydanticVersion:
+    """
+    Get the value of the global variable _pydantic_version.
+    :return: value of the variable
+    """
+    global _pydantic_version
+    return _pydantic_version
 
 def set_custom_template_path(value: Optional[str]) -> None:
     """

src/openapi_python_generator/language_converters/python/generator.py

@@ -31,6 +31,7 @@ def generator(
 
     common.set_use_orjson(use_orjson)
     common.set_custom_template_path(custom_template_path)
+    common.set_pydantic_version(pydantic_version)
 
     if data.components is not None:
         models = generate_models(data.components, pydantic_version)

src/openapi_python_generator/language_converters/python/model_generator.py

@@ -121,10 +121,11 @@ def type_converter(  # noqa: C901
     # With custom string format fields, in order to cast these to strict types (e.g. date, datetime, UUID)
     # orjson is required for JSON serialiation.
     elif (
-            schema.type == "string"
-            and schema.schema_format is not None
-            and schema.schema_format.startswith("uuid")
-            and common.get_use_orjson()
+        schema.type == "string"
+        and schema.schema_format is not None
+        and schema.schema_format.startswith("uuid")
+        # orjson and pydantic v2 both support UUID
+        and (common.get_use_orjson() or common.get_pydantic_version() == PydanticVersion.V2)
     ):
         if len(schema.schema_format) > 4 and schema.schema_format[4].isnumeric():
             uuid_type = schema.schema_format.upper()
@@ -133,13 +134,31 @@ def type_converter(  # noqa: C901
         else:
             converted_type = pre_type + "UUID" + post_type
             import_types = ["from uuid import UUID"]
-    elif schema.type == "string" and schema.schema_format == "date-time" and common.get_use_orjson():
+    elif (
+        schema.type == "string"
+        and schema.schema_format == "date-time"
+        # orjson and pydantic v2 both support datetime
+        and (common.get_use_orjson() or common.get_pydantic_version() == PydanticVersion.V2)
+    ):
         converted_type = pre_type + "datetime" + post_type
         import_types = ["from datetime import datetime"]
-    elif schema.type == "string" and schema.schema_format == "date" and common.get_use_orjson():
+    elif (
+        schema.type == "string"
+        and schema.schema_format == "date"
+        # orjson and pydantic v2 both support date
+        and (common.get_use_orjson() or common.get_pydantic_version() == PydanticVersion.V2)
+    ):
         converted_type = pre_type + "date" + post_type
         import_types = ["from datetime import date"]
-    elif schema.type == "string" and schema.schema_format == "decimal" and common.get_use_orjson():
+    elif (
+        schema.type == "string"
+        and schema.schema_format == "decimal"
+        # orjson does not support Decimal
+        # See https://github.com/ijl/orjson/issues/444
+        and not common.get_use_orjson()
+        # pydantic v2 supports Decimal
+        and common.get_pydantic_version() == PydanticVersion.V2
+    ):
         converted_type = pre_type + "Decimal" + post_type
         import_types = ["from decimal import Decimal"]
     elif schema.type == "string":

src/openapi_python_generator/language_converters/python/service_generator.py

@@ -26,13 +26,29 @@
 
 HTTP_OPERATIONS = ["get", "post", "put", "delete", "options", "head", "patch", "trace"]
 
+def _generate_body_dump_expression(data = "data") -> str:
+    """
+    Generate expression for dumping abstract body as a dictionary.
+    """
+
+    # Use old v1 method for pydantic
+    if common.get_pydantic_version() == common.PydanticVersion.V1:
+        return f"{data}.dict()"
+
+    # Dump model but allow orjson to serialise (fastest)
+    if common.get_use_orjson():
+        return f"{data}.model_dump()"
+
+    # rely on pydantic v2 to serialise (slowest, but best compatibility)
+    return f"{data}.model_dump_json()"
+
 
 def generate_body_param(operation: Operation) -> Union[str, None]:
     if operation.requestBody is None:
         return None
     else:
         if isinstance(operation.requestBody, Reference):
-            return "data.dict()"
+            return _generate_body_dump_expression("data")
 
         if operation.requestBody.content is None:
             return None  # pragma: no cover
@@ -46,11 +62,12 @@ def generate_body_param(operation: Operation) -> Union[str, None]:
             return None  # pragma: no cover
 
         if isinstance(media_type.media_type_schema, Reference):
-            return "data.dict()"
+            return _generate_body_dump_expression("data")
         elif isinstance(media_type.media_type_schema, Schema):
             schema = media_type.media_type_schema
             if schema.type == "array":
-                return "[i.dict() for i in data]"
+                expression = _generate_body_dump_expression("i")
+                return f"[{expression} for i in data]"
             elif schema.type == "object":
                 return "data"
             else:

tests/build_test_api/api.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, date
 from typing import List
 from typing import Optional
 
@@ -20,17 +20,19 @@ class User(BaseModel):
     username: str
     email: str
     password: str
-    is_active: Optional[bool]
+    is_active: Optional[bool] = None
+    created_at: Optional[datetime] = None
+    birthdate: Optional[date] = None
 
 
 class Team(BaseModel):
     id: int
     name: str
     description: str
-    is_active: Optional[bool]
-    created_at: Optional[datetime]
-    updated_at: Optional[datetime]
-    users: Optional[List[User]]
+    is_active: Optional[bool] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+    users: Optional[List[User]] = None
 
 
 @app.get("/", response_model=RootResponse, tags=["general"])

tests/conftest.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Dict
 from typing import Generator
+from openapi_python_generator.language_converters.python import common
 
 import pytest
 from openapi_pydantic.v3.v3_0 import OpenAPI
@@ -32,3 +33,52 @@ def model_data_with_cleanup_fixture(model_data) -> OpenAPI:  # type: ignore
     if test_result_path.exists():
         # delete folder and all subfolders
         shutil.rmtree(test_result_path)
+
+
+@pytest.fixture
+def with_orjson_enabled():
+    """
+    Fixture to enable orjson for the duration of the test
+    """
+    orjson_usage = common.get_use_orjson()
+    common.set_use_orjson(True)
+    try:
+        yield
+    finally:
+        common.set_use_orjson(orjson_usage)
+
+@pytest.fixture
+def with_orjson_disabled():
+    """
+    Fixture to enable orjson for the duration of the test
+    """
+    orjson_usage = common.get_use_orjson()
+    common.set_use_orjson(False)
+    try:
+        yield
+    finally:
+        common.set_use_orjson(orjson_usage)
+
+@pytest.fixture
+def with_pydantic_v1():
+    """
+    Fixture to set pydantic to v1 for the duration of the test
+    """
+    pydantic_version = common.get_pydantic_version()
+    common.set_pydantic_version(common.PydanticVersion.V1)
+    try:
+        yield
+    finally:
+        common.set_pydantic_version(pydantic_version)
+
+@pytest.fixture
+def with_pydantic_v2():
+    """
+    Fixture to set pydantic to v2 for the duration of the test
+    """
+    pydantic_version = common.get_pydantic_version()
+    common.set_pydantic_version(common.PydanticVersion.V2)
+    try:
+        yield
+    finally:
+        common.set_pydantic_version(pydantic_version)

tests/test_data/test_api.json

@@ -532,11 +532,6 @@
             "title": "Birthdate",
             "type": "string",
             "format": "date"
-          },
-          "position": {
-            "title": "Position",
-            "type": "string",
-            "format": "decimal"
           }
         }
       },

tests/test_generated_code.py

@@ -64,7 +64,7 @@ def test_set_auth_token():
     ],
 )
 @respx.mock
-def test_generate_code(model_data_with_cleanup, library, use_orjson, custom_ip):
+def test_generate_code(model_data_with_cleanup, library, use_orjson, custom_ip, with_pydantic_v2):
     generate_data(test_data_path, test_result_path, library, use_orjson=use_orjson)
     result = generator(model_data_with_cleanup, library_config_dict[library])
 
@@ -356,8 +356,8 @@ def test_generate_code(model_data_with_cleanup, library, use_orjson, custom_ip):
         name="team1",
         description="team1",
         is_active=True,
-        created_at="",
-        updated_at="",
+        created_at=None,
+        updated_at=None,
     )
 
     exec_code_base = f"from .test_result.services.general_service import *\nfrom datetime import datetime\nresp_result = create_team_teams_post(Team(**{data}), passed_api_config)"