From b800dfb0e0cfeda0ed45742ebfa2a91d2bae2ba3 Mon Sep 17 00:00:00 2001
From: sydney-runkle <54324534+sydney-runkle@users.noreply.github.com>
Date: Sun, 9 Jul 2023 08:12:02 -0500
Subject: [PATCH] Fixing Pydantic Extension (#1021)

* fixing pydantic extension

* Update model_schema import source
---
 docs/blueprints/pydantic.py | 21 +++++++--------------
 1 file changed, 7 insertions(+), 14 deletions(-)

diff --git a/docs/blueprints/pydantic.py b/docs/blueprints/pydantic.py
index aff62fa8..d815a0ad 100644
--- a/docs/blueprints/pydantic.py
+++ b/docs/blueprints/pydantic.py
@@ -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