feat: Added pydantic v2 to models

Author: MarcoMuellner

src/openapi_python_generator/__main__.py

@@ -1,12 +1,12 @@
 from typing import Optional
+from enum import Enum
 
 import click
 
 from openapi_python_generator import __version__
-from openapi_python_generator.common import HTTPLibrary
+from openapi_python_generator.common import HTTPLibrary, PydanticVersion
 from openapi_python_generator.generate_data import generate_data
 
-
 @click.command()
 @click.argument("source")
 @click.argument("output")
@@ -38,6 +38,13 @@
     default=None,
     help="Custom template path to use. Allows overriding of the built in templates",
 )
+@click.option(
+    "--pydantic-version",
+    type=click.Choice(["v1", "v2"]),
+    default="v2",
+    show_default=True,
+    help="Pydantic version to use for generated models.",
+)
 @click.version_option(version=__version__)
 def main(
     source: str,
@@ -46,6 +53,7 @@ def main(
     env_token_name: Optional[str] = None,
     use_orjson: bool = False,
     custom_template_path: Optional[str] = None,
+    pydantic_version: PydanticVersion = PydanticVersion.V2,
 ) -> None:
     """
     Generate Python code from an OpenAPI 3.0 specification.
@@ -54,9 +62,9 @@ def main(
     an OUTPUT path, where the resulting client is created.
     """
     generate_data(
-        source, output, library, env_token_name, use_orjson, custom_template_path
+        source, output, library, env_token_name, use_orjson, custom_template_path,pydantic_version
     )
 
 
 if __name__ == "__main__":  # pragma: no cover
-    main()
+    main()

src/openapi_python_generator/common.py

@@ -14,6 +14,10 @@ class HTTPLibrary(str, Enum):
     requests = "requests"
     aiohttp = "aiohttp"
 
+class PydanticVersion(str, Enum):
+    V1 = "v1"
+    V2 = "v2"
+
 
 library_config_dict: Dict[Optional[HTTPLibrary], LibraryConfig] = {
     HTTPLibrary.httpx: LibraryConfig(

src/openapi_python_generator/generate_data.py

@@ -13,7 +13,7 @@
 from openapi_pydantic.v3.v3_0 import OpenAPI
 from pydantic import ValidationError
 
-from .common import HTTPLibrary
+from .common import HTTPLibrary, PydanticVersion
 from .common import library_config_dict
 from .language_converters.python.generator import generator
 from .language_converters.python.jinja_config import SERVICE_TEMPLATE
@@ -140,6 +140,7 @@ def generate_data(
     env_token_name: Optional[str] = None,
     use_orjson: bool = False,
     custom_template_path: Optional[str] = None,
+    pydantic_version: PydanticVersion = PydanticVersion.V2,
 ) -> None:
     """
     Generate Python code from an OpenAPI 3.0 specification.
@@ -153,6 +154,7 @@ def generate_data(
         env_token_name,
         use_orjson,
         custom_template_path,
+        pydantic_version,
     )
 
     write_data(result, output)

src/openapi_python_generator/language_converters/python/api_config_generator.py

@@ -2,8 +2,9 @@
 
 from openapi_pydantic.v3.v3_0 import OpenAPI
 
+from openapi_python_generator.common import PydanticVersion
 from openapi_python_generator.language_converters.python.jinja_config import (
-    API_CONFIG_TEMPLATE,
+    API_CONFIG_TEMPLATE, API_CONFIG_TEMPLATE_PYDANTIC_V2,
 )
 from openapi_python_generator.language_converters.python.jinja_config import (
     create_jinja_env,
@@ -12,15 +13,18 @@
 
 
 def generate_api_config(
-    data: OpenAPI, env_token_name: Optional[str] = None
+        data: OpenAPI, env_token_name: Optional[str] = None,
+        pydantic_version: PydanticVersion = PydanticVersion.V2,
 ) -> APIConfig:
     """
     Generate the API model.
     """
+
+    template_name = API_CONFIG_TEMPLATE_PYDANTIC_V2 if pydantic_version == PydanticVersion.V2 else API_CONFIG_TEMPLATE
     jinja_env = create_jinja_env()
     return APIConfig(
         file_name="api_config",
-        content=jinja_env.get_template(API_CONFIG_TEMPLATE).render(
+        content=jinja_env.get_template(template_name).render(
             env_token_name=env_token_name, **data.dict()
         ),
         base_url=data.servers[0].url if len(data.servers) > 0 else "NO SERVER",

src/openapi_python_generator/language_converters/python/generator.py

@@ -2,6 +2,7 @@
 
 from openapi_pydantic.v3.v3_0 import OpenAPI
 
+from openapi_python_generator.common import PydanticVersion
 from openapi_python_generator.language_converters.python import common
 from openapi_python_generator.language_converters.python.api_config_generator import (
     generate_api_config,
@@ -22,6 +23,7 @@ def generator(
     env_token_name: Optional[str] = None,
     use_orjson: bool = False,
     custom_template_path: Optional[str] = None,
+    pydantic_version: PydanticVersion = PydanticVersion.V2,
 ) -> ConversionResult:
     """
     Generate Python code from an OpenAPI 3.0 specification.
@@ -31,7 +33,7 @@ def generator(
     common.set_custom_template_path(custom_template_path)
 
     if data.components is not None:
-        models = generate_models(data.components)
+        models = generate_models(data.components, pydantic_version)
     else:
         models = []
 
@@ -40,7 +42,7 @@ def generator(
     else:
         services = []
 
-    api_config = generate_api_config(data, env_token_name)
+    api_config = generate_api_config(data, env_token_name, pydantic_version)
 
     return ConversionResult(
         models=models,

src/openapi_python_generator/language_converters/python/jinja_config.py

@@ -9,9 +9,11 @@
 
 ENUM_TEMPLATE = "enum.jinja2"
 MODELS_TEMPLATE = "models.jinja2"
+MODELS_TEMPLATE_PYDANTIC_V2 = "models_pydantic_2.jinja2"
 SERVICE_TEMPLATE = "service.jinja2"
 HTTPX_TEMPLATE = "httpx.jinja2"
 API_CONFIG_TEMPLATE = "apiconfig.jinja2"
+API_CONFIG_TEMPLATE_PYDANTIC_V2 = "apiconfig_pydantic_2.jinja2"
 TEMPLATE_PATH = Path(__file__).parent / "templates"
 
 

src/openapi_python_generator/language_converters/python/model_generator.py

@@ -6,9 +6,10 @@
 import click
 from openapi_pydantic.v3.v3_0 import Schema, Reference, Components
 
+from openapi_python_generator.common import PydanticVersion
 from openapi_python_generator.language_converters.python import common
 from openapi_python_generator.language_converters.python.jinja_config import (
-    ENUM_TEMPLATE,
+    ENUM_TEMPLATE, MODELS_TEMPLATE_PYDANTIC_V2,
 )
 from openapi_python_generator.language_converters.python.jinja_config import (
     MODELS_TEMPLATE,
@@ -22,9 +23,9 @@
 
 
 def type_converter(  # noqa: C901
-    schema: Schema,
-    required: bool = False,
-    model_name: Optional[str] = None,
+        schema: Schema,
+        required: bool = False,
+        model_name: Optional[str] = None,
 ) -> TypeConversion:
     """
     Converts an OpenAPI type to a Python type.
@@ -69,13 +70,13 @@ def type_converter(  # noqa: C901
                     )
 
         original_type = (
-            "tuple<" + ",".join([i.original_type for i in conversions]) + ">"
+                "tuple<" + ",".join([i.original_type for i in conversions]) + ">"
         )
         if len(conversions) == 1:
             converted_type = conversions[0].converted_type
         else:
             converted_type = (
-                "Tuple[" + ",".join([i.converted_type for i in conversions]) + "]"
+                    "Tuple[" + ",".join([i.converted_type for i in conversions]) + "]"
             )
 
         converted_type = pre_type + converted_type + post_type
@@ -101,14 +102,14 @@ def type_converter(  # noqa: C901
                     )
                 )
         original_type = (
-            "union<" + ",".join([i.original_type for i in conversions]) + ">"
+                "union<" + ",".join([i.original_type for i in conversions]) + ">"
         )
 
         if len(conversions) == 1:
             converted_type = conversions[0].converted_type
         else:
             converted_type = (
-                "Union[" + ",".join([i.converted_type for i in conversions]) + "]"
+                    "Union[" + ",".join([i.converted_type for i in conversions]) + "]"
             )
 
         converted_type = pre_type + converted_type + post_type
@@ -120,13 +121,13 @@ def type_converter(  # noqa: C901
     # We only want to auto convert to datetime if orjson is used throghout the code, otherwise we can not
     # serialize it to JSON.
     elif schema.type == "string" and (
-        schema.schema_format is None or not common.get_use_orjson()
+            schema.schema_format is None or not common.get_use_orjson()
     ):
         converted_type = pre_type + "str" + post_type
     elif (
-        schema.type == "string"
-        and schema.schema_format.startswith("uuid")
-        and common.get_use_orjson()
+            schema.type == "string"
+            and schema.schema_format.startswith("uuid")
+            and common.get_use_orjson()
     ):
         if len(schema.schema_format) > 4 and schema.schema_format[4].isnumeric():
             uuid_type = schema.schema_format.upper()
@@ -154,7 +155,8 @@ def type_converter(  # noqa: C901
             original_type = "array<" + converted_reference.type.original_type + ">"
             retVal += converted_reference.type.converted_type
         elif isinstance(schema.items, Schema):
-            original_type = "array<" + (str(schema.items.type.value) if schema.items.type is not None else "unknown")+ ">"
+            original_type = "array<" + (
+                str(schema.items.type.value) if schema.items.type is not None else "unknown") + ">"
             retVal += type_converter(schema.items, True).converted_type
         else:
             original_type = "array<unknown>"
@@ -178,7 +180,7 @@ def type_converter(  # noqa: C901
 
 
 def _generate_property_from_schema(
-    model_name: str, name: str, schema: Schema, parent_schema: Optional[Schema] = None
+        model_name: str, name: str, schema: Schema, parent_schema: Optional[Schema] = None
 ) -> Property:
     """
     Generates a property from a schema. It takes the type of the schema and converts it to a python type, and then
@@ -190,9 +192,9 @@ def _generate_property_from_schema(
     :return: Property
     """
     required = (
-        parent_schema is not None
-        and parent_schema.required is not None
-        and name in parent_schema.required
+            parent_schema is not None
+            and parent_schema.required is not None
+            and name in parent_schema.required
     )
 
     import_type = None
@@ -209,11 +211,11 @@ def _generate_property_from_schema(
 
 
 def _generate_property_from_reference(
-    model_name: str,
-    name: str,
-    reference: Reference,
-    parent_schema: Optional[Schema] = None,
-    force_required: bool = False,
+        model_name: str,
+        name: str,
+        reference: Reference,
+        parent_schema: Optional[Schema] = None,
+        force_required: bool = False,
 ) -> Property:
     """
     Generates a property from a reference. It takes the name of the reference as the type, and then
@@ -225,10 +227,10 @@ def _generate_property_from_reference(
     :return: Property and model to be imported by the file
     """
     required = (
-        parent_schema is not None
-        and parent_schema.required is not None
-        and name in parent_schema.required
-    ) or force_required
+                       parent_schema is not None
+                       and parent_schema.required is not None
+                       and name in parent_schema.required
+               ) or force_required
     import_model = common.normalize_symbol(reference.ref.split("/")[-1])
 
     if import_model == model_name:
@@ -256,13 +258,14 @@ def _generate_property_from_reference(
     )
 
 
-def generate_models(components: Components) -> List[Model]:
+def generate_models(components: Components, pydantic_version: PydanticVersion = PydanticVersion.V2) -> List[Model]:
     """
     Receives components from an OpenAPI 3.0 specification and generates the models from it.
     It does so, by iterating over the components.schemas dictionary. For each schema, it checks if
     it is a normal schema (i.e. simple type like string, integer, etc.), a reference to another schema, or
     an array of types/references. It then computes pydantic models from it using jinja2
     :param components: The components from an OpenAPI 3.0 specification.
+    :param pydantic_version: The version of pydantic to use.
     :return: A list of models.
     """
     models: List[Model] = []
@@ -313,7 +316,9 @@ def generate_models(components: Components) -> List[Model]:
                 )
             properties.append(conv_property)
 
-        generated_content = jinja_env.get_template(MODELS_TEMPLATE).render(
+        template_name = MODELS_TEMPLATE_PYDANTIC_V2 if pydantic_version == PydanticVersion.V2 else MODELS_TEMPLATE
+
+        generated_content = jinja_env.get_template(template_name).render(
             schema_name=name, schema=schema_or_reference, properties=properties
         )
 

src/openapi_python_generator/language_converters/python/templates/apiconfig_pydantic_2.jinja2

@@ -0,0 +1,42 @@
+{% if env_token_name is not none %}import os{% endif %}
+
+from pydantic import BaseModel, Field
+from typing import Optional, Union
+
+class APIConfig(BaseModel):
+    model_config = {
+        "validate_assignment": True
+    }
+
+    base_path: str = {% if servers|length > 0 %} '{{ servers[0].url }}' {% else %} 'NO SERVER' {% endif %}
+
+    verify: Union[bool, str] = True
+{% if env_token_name is none %}
+    access_token : Optional[str] = None
+{% endif %}
+
+    def get_access_token(self) -> Optional[str]:
+{% if env_token_name is not none %}
+        try:
+            return os.environ['{{ env_token_name }}']
+        except KeyError:
+            return None
+{% else %}
+        return self.access_token
+{% endif %}
+
+    def set_access_token(self, value : str):
+{% if env_token_name is not none %}
+        raise Exception("This client was generated with an environment variable for the access token. Please set the environment variable '{{ env_token_name }}' to the access token.")
+{% else %}
+        self.access_token = value
+{% endif %}
+
+class HTTPException(Exception):
+    def __init__(self, status_code: int, message: str):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(f"{status_code} {message}")
+
+    def __str__(self):
+        return f"{self.status_code} {self.message}"