51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
from drf_spectacular.drainage import set_override, warn
|
|
from drf_spectacular.extensions import OpenApiSerializerExtension
|
|
from drf_spectacular.plumbing import ResolvedComponent, build_basic_type
|
|
from drf_spectacular.types import OpenApiTypes
|
|
|
|
|
|
class PydanticExtension(OpenApiSerializerExtension):
|
|
"""
|
|
Allows using pydantic models on @extend_schema(request=..., response=...) to
|
|
describe your API.
|
|
|
|
We only have partial support for pydantic's version of dataclass, due to the way they
|
|
are designed. The outermost class (the @extend_schema argument) has to be a subclass
|
|
of pydantic.BaseModel. Inside this outermost BaseModel, any combination of dataclass
|
|
and BaseModel can be used.
|
|
"""
|
|
|
|
target_class = "pydantic.BaseModel"
|
|
match_subclasses = True
|
|
|
|
def get_name(self, auto_schema, direction):
|
|
# due to the fact that it is complicated to pull out every field member BaseModel class
|
|
# of the entry model, we simply use the class name as string for object. This hack may
|
|
# create false positive warnings, so turn it off. However, this may suppress correct
|
|
# warnings involving the entry class.
|
|
# TODO suppression may be migrated to new ComponentIdentity system
|
|
set_override(self.target, 'suppress_collision_warning', True)
|
|
return self.target.__name__
|
|
|
|
def map_serializer(self, auto_schema, direction):
|
|
# let pydantic generate a JSON schema
|
|
try:
|
|
from pydantic.json_schema import model_json_schema
|
|
except ImportError:
|
|
warn("Only pydantic >= 2 is supported. defaulting to generic object.")
|
|
return build_basic_type(OpenApiTypes.OBJECT)
|
|
|
|
schema = model_json_schema(self.target, ref_template="#/components/schemas/{model}", mode="serialization")
|
|
|
|
# pull out potential sub-schemas and put them into component section
|
|
for sub_name, sub_schema in schema.pop("$defs", {}).items():
|
|
component = ResolvedComponent(
|
|
name=sub_name,
|
|
type=ResolvedComponent.SCHEMA,
|
|
object=sub_name,
|
|
schema=sub_schema,
|
|
)
|
|
auto_schema.registry.register_on_missing(component)
|
|
|
|
return schema
|