Add v3.v3_0_3.util for users who still wants to use v3.0.3 OpenAPI with PydanticSchema

Author: kuimono

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Change Log
 
+## v1.2.1 - 2021-09-10
+
+### Added
+- Add `v3.v3_0_3.util` for users who still wants to use v3.0.3 OpenAPI with `PydanticSchema`
+
 ## v1.2.0 - 2021-06-28
 
 ### Added

README.md

@@ -277,6 +277,16 @@ Please refer to the following for more info:
 | 3.1.0 | [here](https://github.com/kuimono/openapi-schema-pydantic/blob/master/openapi_schema_pydantic/v3/v3_1_0/README.md#non-pydantic-schema-types) |
 | 3.0.3 | [here](https://github.com/kuimono/openapi-schema-pydantic/blob/master/openapi_schema_pydantic/v3/v3_0_3/README.md#non-pydantic-schema-types) |
 
+### Use OpenAPI 3.0.3 instead of 3.1.0
+
+Some UI renderings (e.g. Swagger) still does not support OpenAPI 3.1.0,
+so user may still use the 3.0.3 version by importing from different paths:
+
+```python
+from openapi_schema_pydantic.v3.v3_0_3 import OpenAPI, ...
+from openapi_schema_pydantic.v3.v3_0_3.util import PydanticSchema, construct_open_api_with_schema_class
+```
+
 ## License
 
 [MIT License](https://github.com/kuimono/openapi-schema-pydantic/blob/master/LICENSE)

openapi_schema_pydantic/v3/v3_0_3/util.py

@@ -0,0 +1,122 @@
+import logging
+from typing import Any, List, Set, Type, TypeVar
+
+from pydantic import BaseModel
+from pydantic.schema import schema
+
+from . import Components, OpenAPI, Reference, Schema
+
+PydanticType = TypeVar("PydanticType", bound=BaseModel)
+ref_prefix = "#/components/schemas/"
+
+
+class PydanticSchema(Schema):
+    """Special `Schema` class to indicate a reference from pydantic class"""
+
+    schema_class: Type[PydanticType] = ...
+    """the class that is used for generate the schema"""
+
+
+def construct_open_api_with_schema_class(
+    open_api: OpenAPI,
+    schema_classes: List[Type[PydanticType]] = None,
+    scan_for_pydantic_schema_reference: bool = True,
+    by_alias: bool = True,
+) -> OpenAPI:
+    """
+    Construct a new OpenAPI object, with the use of pydantic classes to produce JSON schemas
+
+    :param open_api: the base `OpenAPI` object
+    :param schema_classes: pydanitic 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
+    :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 scan_for_pydantic_schema_reference:
+        extracted_schema_classes = _handle_pydantic_schema(new_open_api)
+        if schema_classes:
+            schema_classes = list({*schema_classes, *_handle_pydantic_schema(new_open_api)})
+        else:
+            schema_classes = extracted_schema_classes
+
+    if not schema_classes:
+        return open_api
+
+    schema_classes.sort(key=lambda x: x.__name__)
+    logging.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 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.get("definitions"):
+                logging.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.get("definitions").items()}
+        )
+    else:
+        new_open_api.components.schemas = {
+            key: Schema.parse_obj(schema_dict) for key, schema_dict in schema_definitions.get("definitions").items()
+        }
+    return new_open_api
+
+
+def _handle_pydantic_schema(open_api: OpenAPI) -> List[Type[PydanticType]]:
+    """
+    This function traverses the `OpenAPI` object and
+
+    1. Replaces the `PydanticSchema` object with `Reference` object, with correct ref value;
+    2. Extracts the involved schema class from `PydanticSchema` object.
+
+    **This function will mutate the input `OpenAPI` object.**
+
+    :param open_api: the `OpenAPI` object to be traversed and mutated
+    :return: a list of schema classes extracted from `PydanticSchema` objects
+    """
+
+    pydantic_types: Set[Type[PydanticType]] = set()
+
+    def _traverse(obj: Any):
+        if isinstance(obj, BaseModel):
+            fields = obj.__fields_set__
+            for field in fields:
+                child_obj = obj.__getattribute__(field)
+                if isinstance(child_obj, PydanticSchema):
+                    logging.debug(f"PydanticSchema found in {obj.__repr_name__()}: {child_obj}")
+                    obj.__setattr__(field, _construct_ref_obj(child_obj))
+                    pydantic_types.add(child_obj.schema_class)
+                else:
+                    _traverse(child_obj)
+        elif isinstance(obj, list):
+            for index, elem in enumerate(obj):
+                if isinstance(elem, PydanticSchema):
+                    logging.debug(f"PydanticSchema found in list: {elem}")
+                    obj[index] = _construct_ref_obj(elem)
+                    pydantic_types.add(elem.schema_class)
+                else:
+                    _traverse(elem)
+        elif isinstance(obj, dict):
+            for key, value in obj.items():
+                if isinstance(value, PydanticSchema):
+                    logging.debug(f"PydanticSchema found in dict: {value}")
+                    obj[key] = _construct_ref_obj(value)
+                    pydantic_types.add(value.schema_class)
+                else:
+                    _traverse(value)
+
+    _traverse(open_api)
+    return list(pydantic_types)
+
+
+def _construct_ref_obj(pydantic_schema: PydanticSchema):
+    ref_obj = Reference(ref=ref_prefix + pydantic_schema.schema_class.__name__)
+    logging.debug(f"ref_obj={ref_obj}")
+    return ref_obj

tests/v3_0_3/__init__.py

@@ -0,0 +1,142 @@
+import logging
+
+from pydantic import BaseModel, Field
+
+from openapi_schema_pydantic.v3.v3_0_3 import Info, MediaType, OpenAPI, Operation, PathItem, Reference, RequestBody, Response
+from openapi_schema_pydantic.v3.v3_0_3.util import PydanticSchema, construct_open_api_with_schema_class
+
+
+def test_construct_open_api_with_schema_class_1():
+    open_api = construct_base_open_api_1()
+    result_open_api_1 = construct_open_api_with_schema_class(open_api)
+    result_open_api_2 = construct_open_api_with_schema_class(open_api, [PingRequest, PingResponse])
+    assert result_open_api_1.components == result_open_api_2.components
+    assert result_open_api_1 == result_open_api_2
+
+    open_api_json = result_open_api_1.json(by_alias=True, exclude_none=True, indent=2)
+    logging.debug(open_api_json)
+
+
+def test_construct_open_api_with_schema_class_2():
+    open_api_1 = construct_base_open_api_1()
+    open_api_2 = construct_base_open_api_2()
+    result_open_api_1 = construct_open_api_with_schema_class(open_api_1)
+    result_open_api_2 = construct_open_api_with_schema_class(open_api_2, [PingRequest, PingResponse])
+    assert result_open_api_1 == result_open_api_2
+
+
+def test_construct_open_api_with_schema_class_3():
+    open_api_3 = construct_base_open_api_3()
+
+    result_with_alias_1 = construct_open_api_with_schema_class(open_api_3)
+    schema_with_alias = result_with_alias_1.components.schemas["PongResponse"]
+    assert "pong_foo" in schema_with_alias.properties
+    assert "pong_bar" in schema_with_alias.properties
+
+    result_with_alias_2 = construct_open_api_with_schema_class(open_api_3, by_alias=True)
+    assert result_with_alias_1 == result_with_alias_2
+
+    result_without_alias = construct_open_api_with_schema_class(open_api_3, by_alias=False)
+    schema_without_alias = result_without_alias.components.schemas["PongResponse"]
+    assert "resp_foo" in schema_without_alias.properties
+    assert "resp_bar" in schema_without_alias.properties
+
+
+def construct_base_open_api_1() -> OpenAPI:
+    return OpenAPI.parse_obj(
+        {
+            "info": {"title": "My own API", "version": "v0.0.1"},
+            "paths": {
+                "/ping": {
+                    "post": {
+                        "requestBody": {
+                            "content": {"application/json": {"schema": PydanticSchema(schema_class=PingRequest)}}
+                        },
+                        "responses": {
+                            "200": {
+                                "description": "pong",
+                                "content": {"application/json": {"schema": PydanticSchema(schema_class=PingResponse)}},
+                            }
+                        },
+                    }
+                }
+            },
+        }
+    )
+
+
+def construct_base_open_api_2() -> OpenAPI:
+    return OpenAPI(
+        info=Info(title="My own API", version="v0.0.1"),
+        paths={
+            "/ping": PathItem(
+                post=Operation(
+                    requestBody=RequestBody(
+                        content={
+                            "application/json": MediaType(
+                                media_type_schema=Reference(ref="#/components/schemas/PingRequest")
+                            )
+                        }
+                    ),
+                    responses={
+                        "200": Response(
+                            description="pong",
+                            content={
+                                "application/json": MediaType(
+                                    media_type_schema=Reference(ref="#/components/schemas/PingResponse")
+                                )
+                            },
+                        )
+                    },
+                )
+            )
+        },
+    )
+
+
+def construct_base_open_api_3() -> OpenAPI:
+    return OpenAPI(
+        info=Info(title="My own API", version="v0.0.1",),
+        paths={
+            "/ping": PathItem(
+                post=Operation(
+                    requestBody=RequestBody(
+                        content={
+                            "application/json": MediaType(media_type_schema=PydanticSchema(schema_class=PingRequest))
+                        }
+                    ),
+                    responses={
+                        "200": Response(
+                            description="pong",
+                            content={
+                                "application/json": MediaType(
+                                    media_type_schema=PydanticSchema(schema_class=PongResponse)
+                                )
+                            },
+                        )
+                    },
+                )
+            )
+        },
+    )
+
+
+class PingRequest(BaseModel):
+    """Ping Request"""
+
+    req_foo: str = Field(description="foo value of the request")
+    req_bar: str = Field(description="bar value of the request")
+
+
+class PingResponse(BaseModel):
+    """Ping response"""
+
+    resp_foo: str = Field(description="foo value of the response")
+    resp_bar: str = Field(description="bar value of the response")
+
+
+class PongResponse(BaseModel):
+    """Pong response"""
+
+    resp_foo: str = Field(alias="pong_foo", description="foo value of the response")
+    resp_bar: str = Field(alias="pong_bar", description="bar value of the response")