Support Pydantic versions 1 and 2

Author: hathawsh

README.md

@@ -18,6 +18,7 @@ The naming of the classes follows the schema in
 
 ```python
 from openapi_pydantic import OpenAPI, Info, PathItem, Operation, Response
+from openapi_pydantic.compat import PYDANTIC_V2
 
 # Construct OpenAPI by pydantic objects
 open_api = OpenAPI(
@@ -37,7 +38,10 @@ open_api = OpenAPI(
         )
     },
 )
-print(open_api.json(by_alias=True, exclude_none=True, indent=2))
+if PYDANTIC_V2:
+    print(open_api.model_dump_json(by_alias=True, exclude_none=True, indent=2))
+else:
+    print(open_api.json(by_alias=True, exclude_none=True, indent=2))
 ```
 
 Result:
@@ -71,12 +75,13 @@ Result:
 
 ## Take advantage of Pydantic
 
-Pydantic is a great tool, allow you to use object / dict / mixed data for for input.
+Pydantic is a great tool. It allows you to use object / dict / mixed data for input.
 
 The following examples give the same OpenAPI result as above:
 
 ```python
 from openapi_pydantic import parse_obj, OpenAPI, PathItem, Response
+from openapi_pydantic.compat import PYDANTIC_V2
 
 # Construct OpenAPI from dict, inferring the correct schema version
 open_api = parse_obj({
@@ -90,7 +95,8 @@ open_api = parse_obj({
 
 
 # Construct OpenAPI v3.1.0 schema from dict
-open_api = OpenAPI.parse_obj({
+openapi_validate = OpenAPI.model_validate if PYDANTIC_V2 else OpenAPI.parse_obj
+open_api = openapi_validate({
     "info": {"title": "My own API", "version": "v0.0.1"},
     "paths": {
         "/ping": {
@@ -100,7 +106,8 @@ open_api = OpenAPI.parse_obj({
 })
 
 # Construct OpenAPI with mix of dict/object
-open_api = OpenAPI.parse_obj({
+openapi_validate = OpenAPI.model_validate if PYDANTIC_V2 else OpenAPI.parse_obj
+open_api = openapi_validate({
     "info": {"title": "My own API", "version": "v0.0.1"},
     "paths": {
         "/ping": PathItem(
@@ -113,10 +120,10 @@ open_api = OpenAPI.parse_obj({
 ## Use Pydantic classes as schema
 
 - The [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.3.md#schemaObject)
-  in OpenAPI has definitions and tweaks in JSON Schema, which is hard to comprehend and define a good data class
-- Pydantic already has a good way to [create JSON schema](https://pydantic-docs.helpmanual.io/usage/schema/),
-  let's not re-invent the wheel
-  
+  in OpenAPI has definitions and tweaks in JSON Schema, which are hard to comprehend and define a good data class
+- Pydantic already has a good way to [create JSON schema](https://pydantic-docs.helpmanual.io/usage/schema/).
+  Let's not reinvent the wheel.
+
 The approach to deal with this:
 
 1. Use `PydanticSchema` objects to represent the `Schema` in `OpenAPI` object
@@ -126,10 +133,12 @@ The approach to deal with this:
 from pydantic import BaseModel, Field
 
 from openapi_pydantic import OpenAPI
+from openapi_pydantic.compat import PYDANTIC_V2
 from openapi_pydantic.util import PydanticSchema, construct_open_api_with_schema_class
 
 def construct_base_open_api() -> OpenAPI:
-    return OpenAPI.parse_obj({
+    openapi_validate = OpenAPI.model_validate if PYDANTIC_V2 else OpenAPI.parse_obj
+    return openapi_validate({
         "info": {"title": "My own API", "version": "v0.0.1"},
         "paths": {
             "/ping": {
@@ -162,7 +171,10 @@ open_api = construct_base_open_api()
 open_api = construct_open_api_with_schema_class(open_api)
 
 # print the result openapi.json
-print(open_api.json(by_alias=True, exclude_none=True, indent=2))
+if PYDANTIC_V2:
+    print(open_api.model_dump_json(by_alias=True, exclude_none=True, indent=2))
+else:
+    print(open_api.json(by_alias=True, exclude_none=True, indent=2))
 ```
 
 Result:
@@ -259,21 +271,24 @@ Result:
 
 ## Notes
 
-### Use of OpenAPI.json() / OpenAPI.dict()
+### Use of OpenAPI.model_dump() / OpenAPI.model_dump_json() / OpenAPI.json() / OpenAPI.dict()
 
-When using `OpenAPI.json()` / `OpenAPI.dict()` function,
-arguments `by_alias=True, exclude_none=True` has to be in place.
-Otherwise the result json will not fit the OpenAPI standard.
+When using `OpenAPI.model_dump()` / `OpenAPI.model_dump_json()` / `OpenAPI.json()` / `OpenAPI.dict()` functions,
+the arguments `by_alias=True, exclude_none=True` have to be in place.
+Otherwise the resulting json will not fit the OpenAPI standard.
 
 ```python
-# OK
+# OK (Pydantic 2)
+open_api.model_dump_json(by_alias=True, exclude_none=True, indent=2)
+# OK (Pydantic 1)
 open_api.json(by_alias=True, exclude_none=True, indent=2)
 
 # Not good
+open_api.model_dump_json(indent=2)
 open_api.json(indent=2)
 ```
 
-More info about field alias:
+More info about field aliases:
 
 | OpenAPI version | Field alias info |
 | --------------- | ---------------- |
@@ -293,7 +308,7 @@ Please refer to the following for more info:
 ### Use OpenAPI 3.0.3 instead of 3.1.0
 
 Some UI renderings (e.g. Swagger) still do not support OpenAPI 3.1.0.
-It is allowed to use the old 3.0.3 version by importing from different paths:
+The old 3.0.3 version is available by importing from different paths:
 
 ```python
 from openapi_pydantic.v3.v3_0_3 import OpenAPI, ...

openapi_pydantic/compat.py

@@ -0,0 +1,53 @@
+"""Compatibility layer to make this package usable with Pydantic 1 or 2"""
+
+from pydantic.version import VERSION as PYDANTIC_VERSION
+
+PYDANTIC_MAJOR_VERSION = int(PYDANTIC_VERSION.split(".", 1)[0])
+
+if int(PYDANTIC_MAJOR_VERSION) >= 2:
+    PYDANTIC_V2 = True
+else:
+    PYDANTIC_V2 = False
+
+if PYDANTIC_V2:
+    from typing import Literal
+
+    from pydantic import ConfigDict
+    from pydantic.json_schema import JsonSchemaMode, models_json_schema  # type: ignore
+
+    # Pydantic 2 renders JSON schemas using the keyword "$defs"
+    DEFS_KEY = "$defs"
+
+    # Add V1 stubs to this module, but hide them from typing
+    globals().update(
+        {
+            "Extra": None,
+            "v1_schema": None,
+        }
+    )
+
+else:
+    from pydantic import Extra
+    from pydantic.schema import schema as v1_schema
+
+    # Pydantic 1 renders JSON schemas using the keyword "definitions"
+    DEFS_KEY = "definitions"
+
+    # Add V2 stubs to this module, but hide them from typing
+    globals().update(
+        {
+            "ConfigDict": None,
+            "Literal": None,
+            "models_json_schema": None,
+            "JsonSchemaMode": None,
+        }
+    )
+
+__all__ = [
+    "Literal",
+    "ConfigDict",
+    "JsonSchemaMode",
+    "models_json_schema",
+    "Extra",
+    "v1_schema",
+]

openapi_pydantic/util.py

@@ -2,14 +2,22 @@
 from typing import Any, Generic, List, Optional, Set, Type, TypeVar
 
 from pydantic import BaseModel
-from pydantic.schema import schema
+
+from openapi_pydantic.compat import (
+    DEFS_KEY,
+    PYDANTIC_V2,
+    JsonSchemaMode,
+    models_json_schema,
+    v1_schema,
+)
 
 from . import Components, OpenAPI, Reference, Schema
 
 logger = logging.getLogger(__name__)
 
 PydanticType = TypeVar("PydanticType", bound=BaseModel)
 ref_prefix = "#/components/schemas/"
+ref_template = "#/components/schemas/{model}"
 
 
 class PydanticSchema(Schema, Generic[PydanticType]):
@@ -19,28 +27,46 @@ class PydanticSchema(Schema, Generic[PydanticType]):
     """the class that is used for generate the schema"""
 
 
+def get_mode(
+    cls: Type[BaseModel], default: JsonSchemaMode = "validation"
+) -> JsonSchemaMode:
+    """Get the JSON schema mode for a model class.
+
+    The mode can be either "serialization" or "validation". In validation mode,
+    computed fields are dropped and optional fields remain optional. In
+    serialization mode, computed and optional fields are required.
+    """
+    if not hasattr(cls, "model_config"):
+        return default
+    return cls.model_config.get("json_schema_mode", default)
+
+
 def construct_open_api_with_schema_class(
     open_api: OpenAPI,
     schema_classes: Optional[List[Type[BaseModel]]] = None,
     scan_for_pydantic_schema_reference: bool = True,
     by_alias: bool = True,
 ) -> OpenAPI:
     """
-    Construct a new OpenAPI object, utilising pydantic classes to produce JSON schemas
+    Construct a new OpenAPI object, utilising pydantic classes to produce JSON schemas.
 
     :param open_api: the base `OpenAPI` object
-    :param schema_classes: pydanitic classes that their schema will be used
+    :param schema_classes: Pydantic classes that their schema will be used
                            "#/components/schemas" values
     :param scan_for_pydantic_schema_reference: flag to indicate if scanning for
-                                               `PydanticSchemaReference` class is
-                                               needed for "#/components/schemas" value
-                                               updates
+                                               `PydanticSchemaReference` class
+                                               is needed for "#/components/schemas"
+                                               value updates
     :param by_alias: construct schema by alias (default is True)
     :return: new OpenAPI object with "#/components/schemas" values updated.
              If there is no update in "#/components/schemas" values, the original
              `open_api` will be returned.
     """
-    new_open_api: OpenAPI = open_api.copy(deep=True)
+    if PYDANTIC_V2:
+        new_open_api = open_api.model_copy(deep=True)
+    else:
+        new_open_api = open_api.copy(deep=True)
+
     if scan_for_pydantic_schema_reference:
         extracted_schema_classes = _handle_pydantic_schema(new_open_api)
         if schema_classes:
@@ -57,28 +83,37 @@ def construct_open_api_with_schema_class(
     logger.debug(f"schema_classes{schema_classes}")
 
     # update new_open_api with new #/components/schemas
-    schema_definitions = schema(
-        schema_classes, by_alias=by_alias, ref_prefix=ref_prefix
-    )
+    if PYDANTIC_V2:
+        _key_map, schema_definitions = models_json_schema(
+            [(c, get_mode(c)) for c in schema_classes],
+            by_alias=by_alias,
+            ref_template=ref_template,
+        )
+    else:
+        schema_definitions = v1_schema(
+            schema_classes, by_alias=by_alias, ref_prefix=ref_prefix
+        )
+
+    schema_validate = Schema.model_validate if PYDANTIC_V2 else Schema.parse_obj
     if not new_open_api.components:
         new_open_api.components = Components()
     if new_open_api.components.schemas:
         for existing_key in new_open_api.components.schemas:
-            if existing_key in schema_definitions["definitions"]:
+            if existing_key in schema_definitions[DEFS_KEY]:
                 logger.warning(
                     f'"{existing_key}" already exists in {ref_prefix}. '
                     f'The value of "{ref_prefix}{existing_key}" will be overwritten.'
                 )
         new_open_api.components.schemas.update(
             {
-                key: Schema.parse_obj(schema_dict)
-                for key, schema_dict in schema_definitions["definitions"].items()
+                key: schema_validate(schema_dict)
+                for key, schema_dict in schema_definitions[DEFS_KEY].items()
             }
         )
     else:
         new_open_api.components.schemas = {
-            key: Schema.parse_obj(schema_dict)
-            for key, schema_dict in schema_definitions["definitions"].items()
+            key: schema_validate(schema_dict)
+            for key, schema_dict in schema_definitions[DEFS_KEY].items()
         }
     return new_open_api
 
@@ -101,7 +136,7 @@ def _handle_pydantic_schema(open_api: OpenAPI) -> List[Type[BaseModel]]:
 
     def _traverse(obj: Any) -> None:
         if isinstance(obj, BaseModel):
-            fields = obj.__fields_set__
+            fields = obj.model_fields_set if PYDANTIC_V2 else obj.__fields_set__
             for field in fields:
                 child_obj = obj.__getattribute__(field)
                 if isinstance(child_obj, PydanticSchema):

openapi_pydantic/v3/parser.py

@@ -2,14 +2,28 @@
 
 from pydantic import BaseModel, Field
 
+from openapi_pydantic.compat import PYDANTIC_V2
+
 from .v3_0_3 import OpenAPI as OpenAPIv3_0
 from .v3_1_0 import OpenAPI as OpenAPIv3_1
 
+OpenAPIv3 = Union[OpenAPIv3_1, OpenAPIv3_0]
+
+if PYDANTIC_V2:
+    from pydantic import RootModel
+
+    class _OpenAPIV2(RootModel):
+        root: OpenAPIv3 = Field(discriminator="openapi")
+
+else:
 
-class _OpenAPI(BaseModel):
-    __root__: Union[OpenAPIv3_1, OpenAPIv3_0] = Field(discriminator="openapi")
+    class _OpenAPIV1(BaseModel):
+        __root__: OpenAPIv3 = Field(discriminator="openapi")
 
 
-def parse_obj(data: Any) -> Union[OpenAPIv3_1, OpenAPIv3_0]:
+def parse_obj(data: Any) -> OpenAPIv3:
     """Parse a raw object into an OpenAPI model with version inference."""
-    return _OpenAPI.parse_obj(data).__root__
+    if PYDANTIC_V2:
+        return _OpenAPIV2.model_validate(data).root
+    else:
+        return _OpenAPIV1.parse_obj(data).__root__

openapi_pydantic/v3/v3_0_3/README.md

@@ -20,7 +20,7 @@ the following fields are used with [alias](https://pydantic-docs.helpmanual.io/u
 > <a name="header_param_in"></a>The "in" field in Header object is actually a constant (`{"in": "header"}`).
 
 > For convenience of object creation, the classes mentioned in above
-> has configured `allow_population_by_field_name=True`.
+> have configured `allow_population_by_field_name=True` (Pydantic V1) or `populate_by_name=True` (Pydantic V2).
 >
 > Reference: [Pydantic's Model Config](https://pydantic-docs.helpmanual.io/usage/model_config/)
 

openapi_pydantic/v3/v3_0_3/__init__.py

@@ -6,6 +6,8 @@
 https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.3.md#table-of-contents
 """
 
+from openapi_pydantic.compat import PYDANTIC_V2
+
 from .callback import Callback as Callback
 from .components import Components as Components
 from .contact import Contact as Contact
@@ -40,6 +42,12 @@
 from .xml import XML as XML
 
 # resolve forward references
-Encoding.update_forward_refs(Header=Header)
-Schema.update_forward_refs()
-Operation.update_forward_refs(PathItem=PathItem)
+if PYDANTIC_V2:
+    Encoding.model_rebuild()
+    OpenAPI.model_rebuild()
+    Components.model_rebuild()
+    Operation.model_rebuild()
+else:
+    Encoding.update_forward_refs(Header=Header)
+    Schema.update_forward_refs()
+    Operation.update_forward_refs(PathItem=PathItem)