diff --git a/src/apispec_plugins/__init__.py b/src/apispec_plugins/__init__.py index 6130a8e..7d1b23d 100644 --- a/src/apispec_plugins/__init__.py +++ b/src/apispec_plugins/__init__.py @@ -1,12 +1,13 @@ from importlib import metadata from .utils import spec_from -from .ext.pydantic import PydanticPlugin +from .ext.pydantic import BaseModel, PydanticPlugin from .webframeworks.flask import FlaskPlugin from .base.mixin import DataclassSchemaMixin, RegistryMixin __version__ = metadata.version("apispec-plugins") __all__ = ( + BaseModel, DataclassSchemaMixin, FlaskPlugin, PydanticPlugin, diff --git a/src/apispec_plugins/ext/pydantic.py b/src/apispec_plugins/ext/pydantic.py index 93089a5..886530c 100644 --- a/src/apispec_plugins/ext/pydantic.py +++ b/src/apispec_plugins/ext/pydantic.py @@ -5,8 +5,12 @@ from apispec import BasePlugin, APISpec from apispec.exceptions import APISpecError, DuplicateComponentNameError -from apispec_plugins.base.registry import Registry -from pydantic import BaseModel +from apispec_plugins.base.mixin import RegistryMixin +from pydantic import BaseModel as PBaseModel + + +class BaseModel(PBaseModel, RegistryMixin): + """Extend BaseModel with Registry""" class PydanticPlugin(BasePlugin): @@ -152,7 +156,7 @@ def resolve_schema_instance( elif isinstance(schema, BaseModel): return schema.__class__ elif isinstance(schema, str): - return Registry.get_cls(schema) + return BaseModel.get_cls(schema) return None @classmethod diff --git a/tests/conftest.py b/tests/conftest.py index 5535cf9..6951d5b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,8 @@ from typing import Optional -from apispec_plugins.base.mixin import RegistryMixin -from pydantic import BaseModel +from apispec_plugins.ext.pydantic import BaseModel -class Pet(BaseModel, RegistryMixin): +class Pet(BaseModel): id: Optional[int] name: str