Merge pull request #50 from codectl/fix/#47

register HTTPResponse model response if registry being used

Author: codectl

README.rst

@@ -57,7 +57,7 @@ Example Usage
     from typing import Optional
 
     from apispec import APISpec
-    from apispec_plugins.base.registry import RegistryMixin
+    from apispec_plugins.base.mixin import RegistryMixin
     from apispec_plugins.ext.pydantic import PydanticPlugin
     from apispec_plugins.webframeworks.flask import FlaskPlugin
     from flask import Flask
@@ -68,7 +68,7 @@ Example Usage
     spec = APISpec(
         title="Pet Store",
         version="1.0.0",
-        openapi_version="3.1.0",
+        openapi_version="3.0.3",
         info=dict(description="A minimal pet store API"),
         plugins=(FlaskPlugin(), PydanticPlugin()),
     )

poetry.lock

@@ -43,15 +43,16 @@ classifiers = [
 [tool.poetry.dependencies]
 apispec = { extras = ["yaml"], version = "^6.3.0" }
 flask = { version = "^2.1.3", optional = true }
-pydantic = { version = "^1.10.6", optional = true }
+pydantic = { version = "^1.10.7", optional = true }
 python = "^3.8"
 
 [tool.poetry.dev-dependencies]
 coverage = "^7.2.2"
 flask = "^2.1.2"
 pre-commit = "^3.1.1"
-pydantic = "^1.10.5"
+pydantic = "^1.10.7"
 pytest = "^7.2.2"
+pytest-mock = "^3.10.0"
 
 [tool.poetry.extras]
 flask = ["Flask"]

pyproject.toml

@@ -0,0 +1,44 @@
+from dataclasses import MISSING, fields
+from typing import get_args
+
+from apispec import APISpec
+from apispec.ext.marshmallow.openapi import OpenAPIConverter, marshmallow as ma
+from apispec_plugins.base.registry import Registry
+
+
+class RegistryMixin(Registry):
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.register(cls)
+
+
+class DataclassSchemaMixin:
+    @classmethod
+    def schema(cls):
+
+        # resolve pydantic schema
+        model = getattr(cls, "__pydantic_model__", None)
+        if model:
+            Registry.register(model)
+            return model.schema()
+
+        # or fallback to marshmallow resolver
+        return cls.dataclass_schema()
+
+    @classmethod
+    def dataclass_schema(cls, openapi_version="2.0"):
+        openapi_converter = OpenAPIConverter(
+            openapi_version=openapi_version,
+            schema_name_resolver=lambda f: None,
+            spec=APISpec("", "", openapi_version),
+        )
+
+        def schema_type(t):
+            return ma.Schema.TYPE_MAPPING[next(iter(get_args(t)), t)]
+
+        schema_dict = {
+            f.name: schema_type(f.type)(data_key=f.name, required=f.default is MISSING)
+            for f in fields(cls)
+        }
+        schema = ma.Schema.from_dict(schema_dict)
+        return openapi_converter.schema2jsonschema(schema)

src/apispec_plugins/base/mixin.py

@@ -3,20 +3,20 @@
 from typing import TypeVar
 
 __all__ = (
-    "RegistryMixin",
+    "Registry",
     "RegistryError",
 )
 
 T = TypeVar("T")
 
 
-class RegistryMixin:
-
+class Registry:
     _registry: dict[str, T] = {}
 
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        cls._registry[cls.__name__] = cls
+    @classmethod
+    def register(cls, record: T):
+        if record.__name__ not in cls._registry:
+            cls._registry[record.__name__] = record
 
     @classmethod
     def get_registry(cls):

src/apispec_plugins/base/registry.py

@@ -1,13 +1,17 @@
-from __future__ import annotations
+try:
+    from pydantic.dataclasses import dataclass
+except ImportError:
+    from dataclasses import dataclass
+from typing import Optional
 
-from dataclasses import dataclass
+from apispec_plugins.base.mixin import DataclassSchemaMixin
 
 
 __all__ = (
     "AuthSchemes",
-    "HTTPResponse",
     "Server",
     "Tag",
+    "HTTPResponse",
 )
 
 
@@ -18,19 +22,19 @@ class BasicAuth:
         scheme: str = "basic"
 
 
-@dataclass
-class HTTPResponse:
-    code: int
-    description: str | None = None
-
-
 @dataclass
 class Server:
     url: str
-    description: str | None = None
+    description: Optional[str] = None
 
 
 @dataclass
 class Tag:
     name: str
-    description: str | None = None
+    description: Optional[str] = None
+
+
+@dataclass
+class HTTPResponse(DataclassSchemaMixin):
+    code: int
+    description: Optional[str] = None

src/apispec_plugins/base/types.py

@@ -5,7 +5,7 @@
 
 from apispec import BasePlugin, APISpec
 from apispec.exceptions import APISpecError, DuplicateComponentNameError
-from apispec_plugins.base import registry
+from apispec_plugins.base.registry import Registry
 from pydantic import BaseModel
 
 
@@ -152,7 +152,7 @@ def resolve_schema_instance(
         elif isinstance(schema, BaseModel):
             return schema.__class__
         elif isinstance(schema, str):
-            return registry.RegistryMixin.get_cls(schema)
+            return Registry.get_cls(schema)
         return None
 
     @classmethod

src/apispec_plugins/ext/pydantic.py

@@ -2,19 +2,17 @@
 import re
 import typing
 import urllib.parse
-from collections.abc import Sequence
-from dataclasses import MISSING, asdict, fields
+from dataclasses import asdict
 
 from apispec import yaml_utils
-
 from apispec_plugins.base import types
 
+
 __all__ = (
     "spec_from",
     "load_method_specs",
     "load_specs_from_docstring",
     "path_parser",
-    "dataclass_schema_resolver",
     "base_template",
 )
 
@@ -81,31 +79,6 @@ def path_parser(path, **kwargs):
     return parsed
 
 
-def dataclass_schema_resolver(schema):
-    """A schema resolver for dataclasses."""
-
-    def _resolve_field_type(f):
-        if f.type == str:
-            return "string"
-        if f.type == int:
-            return "integer"
-        if f.type == float:
-            return "number"
-        if f.type == bool:
-            return "boolean"
-        elif isinstance(field.type, Sequence):
-            return "array"
-        return "object"
-
-    definition = {"type": "object", "properties": {}, "required": []}
-    for field in fields(schema):
-        name = field.name
-        definition["properties"][name] = {"type": _resolve_field_type(field)}
-        if field.default == MISSING and field.default_factory == MISSING:
-            definition["required"].append(name)
-    return definition
-
-
 def base_template(
     openapi_version: str,
     info: dict = None,

src/apispec_plugins/utils.py

@@ -77,9 +77,7 @@ def operation_helper(self, path=None, operations=None, **kwargs):
                         if http_schema_name not in self.spec.components.schemas:
                             self.spec.components.schema(
                                 component_id=http_schema_name,
-                                component=spec_utils.dataclass_schema_resolver(
-                                    types.HTTPResponse
-                                ),
+                                component=types.HTTPResponse.schema(),
                             )
 
                         if schema_name not in self.spec.components.responses:

src/apispec_plugins/webframeworks/flask.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+from apispec_plugins.base.mixin import RegistryMixin
+from pydantic import BaseModel
+
+
+class Pet(BaseModel, RegistryMixin):
+    id: Optional[int]
+    name: str