Skip to content

Commit

Permalink
Fixing Pydantic Extension (#1021)
Browse files Browse the repository at this point in the history
* fixing pydantic extension

* Update model_schema import source
  • Loading branch information
sydney-runkle authored Jul 9, 2023
1 parent e0f749e commit b800dfb
Showing 1 changed file with 7 additions and 14 deletions.
21 changes: 7 additions & 14 deletions docs/blueprints/pydantic.py
Original file line number Diff line number Diff line change
@@ -1,36 +1,29 @@
from drf_spectacular.extensions import OpenApiSerializerExtension
from drf_spectacular.plumbing import ResolvedComponent

from pydantic.schema import model_schema


class PydanticExtension(OpenApiSerializerExtension):
target_class = 'pydantic.BaseModel'
target_class = "pydantic.BaseModel"
match_subclasses = True

def get_name(self, auto_schema, direction):
return self.target.__name__

def map_serializer(self, auto_schema, direction):
def translate_refs(obj, key=None):
if isinstance(obj, dict):
return {k: translate_refs(v, k) for k, v in obj.items()}
elif isinstance(obj, list):
return [translate_refs(i) for i in obj]
elif key == '$ref':
return obj.replace('#/definitions/', '#/components/schemas/')
else:
return obj

# let pydantic generate a JSON schema
schema = self.target.schema()
schema = model_schema(self.target, ref_prefix="#/components/schemas/")

# pull out potential sub-schemas and put them into component section
for sub_name, sub_schema in schema.pop('definitions', {}).items():
for sub_name, sub_schema in schema.pop("definitions", {}).items():
component = ResolvedComponent(
name=sub_name,
type=ResolvedComponent.SCHEMA,
object=sub_name,
schema=translate_refs(sub_schema),
schema=sub_schema,
)
auto_schema.registry.register_on_missing(component)

return translate_refs(schema)
return schema

0 comments on commit b800dfb

Please sign in to comment.