Support customizing component name for pydantic models (#12)

* support customizing pydantic model name in component schema

* added constraint to pydantic

* refactor

Author: Goldziher

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
         exclude: "\\.idea/(.)*"
       - id: trailing-whitespace
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.37.3
+    rev: v2.38.0
     hooks:
       - id: pyupgrade
         args: ["--py37-plus"]
@@ -19,7 +19,7 @@ repos:
     hooks:
       - id: isort
   - repo: https://github.com/psf/black
-    rev: 22.6.0
+    rev: 22.8.0
     hooks:
       - id: black
         args: [--config=./pyproject.toml]
@@ -73,11 +73,11 @@ repos:
             "flake8-noqa",
           ]
   - repo: https://github.com/johnfraney/flake8-markdown
-    rev: v0.3.0
+    rev: v0.4.0
     hooks:
       - id: flake8-markdown
   - repo: https://github.com/pycqa/pylint
-    rev: "v2.15.0"
+    rev: "v2.15.2"
     hooks:
       - id: pylint
         exclude: "test_*|docs"

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+[v1.3.0]
+
+- add support for `__schema_name__` dunder attribute on pydantic models
+
 [v1.2.0]
 
 - update to pydantic `v1.1.0`

poetry.lock

@@ -1,6 +1,6 @@
-from typing import TYPE_CHECKING, Any, List, Optional, Set, Type, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Set, Type, TypeVar, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, create_model
 from pydantic.schema import schema
 
 from pydantic_openapi_schema import v3_1_0
@@ -9,6 +9,7 @@
     from typing import Dict
 
 REF_PREFIX = "#/components/schemas/"
+SCHEMA_NAME_ATTRIBUTE = "__schema_name__"
 
 T = TypeVar("T", bound=v3_1_0.OpenAPI)
 
@@ -22,45 +23,35 @@ class OpenAPI310PydanticSchema(v3_1_0.Schema):
 
 def construct_open_api_with_schema_class(
     open_api_schema: T,
-    schema_classes: Optional[List[Type[BaseModel]]] = None,
-    scan_for_pydantic_schema_reference: bool = True,
-    by_alias: bool = True,
 ) -> T:
     """Construct a new OpenAPI object, with the use of pydantic classes to
     produce JSON schemas.
 
     Args:
-        open_api_schema: the base `OpenAPI` object
-        schema_classes: pydantic classes that their schema will be used as "#/components/schemas" values
-        scan_for_pydantic_schema_reference: flag to indicate if scanning for `PydanticSchemaReference`
-            class is needed for "#/components/schemas" value updates
-        by_alias: construct schema by alias (default is True)
+        open_api_schema: An instance of the OpenAPI model.
 
     Returns:
         new OpenAPI object with "#/components/schemas" values updated. If there is no update in
             "#/components/schemas" values, the original `open_api` will be returned.
     """
     copied_schema = open_api_schema.copy(deep=True)
-
-    if scan_for_pydantic_schema_reference:
-        extracted_schema_classes = extract_pydantic_types_to_openapi_components(
-            obj=copied_schema, ref_class=v3_1_0.Reference
-        )
-        schema_classes = list(
-            {*schema_classes, *extracted_schema_classes} if schema_classes else extracted_schema_classes
-        )
+    schema_classes = list(extract_pydantic_types_to_openapi_components(obj=copied_schema, ref_class=v3_1_0.Reference))
 
     if not schema_classes:
         return open_api_schema
 
     if not copied_schema.components:
         copied_schema.components = v3_1_0.Components(schemas={})
-    elif not copied_schema.components.schemas:
+    if copied_schema.components.schemas is None:  # pragma: no cover
         copied_schema.components.schemas = cast("Dict[str, Any]", {})
 
+    schema_classes = [
+        cls if not hasattr(cls, "__schema_name__") else create_model(getattr(cls, SCHEMA_NAME_ATTRIBUTE), __base__=cls)
+        for cls in schema_classes
+    ]
     schema_classes.sort(key=lambda x: x.__name__)
-    schema_definitions = schema(schema_classes, by_alias=by_alias, ref_prefix=REF_PREFIX)["definitions"]
-    copied_schema.components.schemas.update(  # type: ignore
+    schema_definitions = schema(schema_classes, ref_prefix=REF_PREFIX)["definitions"]
+    copied_schema.components.schemas.update(
         {key: v3_1_0.Schema.parse_obj(schema_dict) for key, schema_dict in schema_definitions.items()}
     )
     return copied_schema
@@ -84,22 +75,34 @@ def extract_pydantic_types_to_openapi_components(obj: Any, ref_class: Type[v3_1_
         for field in fields:
             child_obj = getattr(obj, field)
             if isinstance(child_obj, OpenAPI310PydanticSchema):
-                setattr(obj, field, ref_class(ref=REF_PREFIX + child_obj.schema_class.__name__))
+                setattr(obj, field, ref_class(ref=create_ref_prefix(child_obj.schema_class)))
                 pydantic_schemas.add(child_obj.schema_class)
             else:
                 pydantic_schemas.update(extract_pydantic_types_to_openapi_components(child_obj, ref_class=ref_class))
     elif isinstance(obj, list):
         for index, elem in enumerate(obj):
             if isinstance(elem, OpenAPI310PydanticSchema):
-                obj[index] = ref_class(ref=REF_PREFIX + elem.schema_class.__name__)
+                obj[index] = ref_class(ref=create_ref_prefix(elem.schema_class))
                 pydantic_schemas.add(elem.schema_class)
             else:
                 pydantic_schemas.update(extract_pydantic_types_to_openapi_components(elem, ref_class=ref_class))
     elif isinstance(obj, dict):
         for key, value in obj.items():
             if isinstance(value, OpenAPI310PydanticSchema):
-                obj[key] = ref_class(ref=REF_PREFIX + value.schema_class.__name__)
+                obj[key] = ref_class(ref=create_ref_prefix(value.schema_class))
                 pydantic_schemas.add(value.schema_class)
             else:
                 pydantic_schemas.update(extract_pydantic_types_to_openapi_components(value, ref_class=ref_class))
     return pydantic_schemas
+
+
+def create_ref_prefix(model: Type[BaseModel]) -> str:
+    """
+
+    Args:
+        model: Pydantic model instance.
+
+    Returns:
+        A prefixed name.
+    """
+    return REF_PREFIX + getattr(model, SCHEMA_NAME_ATTRIBUTE, model.__name__)

pydantic_openapi_schema/utils/utils.py

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pydantic-openapi-schema"
-version = "1.2.0"
+version = "1.3.0"
 description = "OpenAPI Schema using pydantic. Forked for Starlite-API from 'openapi-schema-pydantic'."
 authors = ["Na'aman Hirschfeld <nhirschfeld@gmail.com>"]
 maintainers = ["Na'aman Hirschfeld <nhirschfeld@gmail.com>", "Peter Schutt <peter.github@proton.me>", "Cody Fincher <cody.fincher@gmail.com>"]
@@ -28,7 +28,7 @@ packages = [
 ]
 [tool.poetry.dependencies]
 python = ">=3.7"
-pydantic = "*"
+pydantic = ">=1.10.0"
 email-validator = "*"
 
 [tool.poetry.dev-dependencies]

pyproject.toml

@@ -1,10 +1,7 @@
-from pydantic import BaseModel, Extra
-from pydantic.schema import schema
-
 from pydantic_openapi_schema.v3_1_0 import Reference, Schema
 
 
-def test_schema() -> None:
+def test_schema_parse_obj() -> None:
     parsed_schema = Schema.parse_obj(
         {
             "title": "reference list",
@@ -16,30 +13,3 @@ def test_schema() -> None:
     assert isinstance(parsed_schema.allOf, list)
     assert isinstance(parsed_schema.allOf[0], Reference)
     assert parsed_schema.allOf[0].ref == "#/definitions/TestType"
-
-
-def test_issue_4() -> None:
-    """https://github.com/kuimono/openapi-schema-pydantic/issues/4."""
-
-    class TestModel(BaseModel):
-        test_field: str
-
-        class Config:
-            extra = Extra.forbid
-
-    schema_definition = schema([TestModel])
-    assert schema_definition == {
-        "definitions": {
-            "TestModel": {
-                "title": "TestModel",
-                "type": "object",
-                "properties": {"test_field": {"title": "Test Field", "type": "string"}},
-                "required": ["test_field"],
-                "additionalProperties": False,
-            }
-        }
-    }
-
-    # allow "additionalProperties" to have boolean value
-    result = Schema.parse_obj(schema_definition["definitions"]["TestModel"])
-    assert result.additionalProperties is False