Fix PolymorphicSerializer type field handling #885 #958

tfranzel · Mar 12, 2023 · 06a39ad · 06a39ad
1 parent fde16e4
commit 06a39ad
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 24 deletions.
diff --git a/drf_spectacular/contrib/rest_polymorphic.py b/drf_spectacular/contrib/rest_polymorphic.py
@@ -1,5 +1,9 @@
 from drf_spectacular.extensions import OpenApiSerializerExtension
-from drf_spectacular.plumbing import warn
+from drf_spectacular.plumbing import (
+    ResolvedComponent, build_basic_type, build_object_type, is_patched_serializer, warn,
+)
+from drf_spectacular.settings import spectacular_settings
+from drf_spectacular.types import OpenApiTypes
 
 
 class PolymorphicSerializerExtension(OpenApiSerializerExtension):
@@ -17,7 +21,13 @@ def map_serializer(self, auto_schema, direction):
             component = auto_schema.resolve_serializer(sub_serializer, direction)
             if not component:
                 continue
-            sub_components.append((resource_type, component.ref))
+            typed_component = self.build_typed_component(
+                auto_schema=auto_schema,
+                component=component,
+                resource_type_field_name=serializer.resource_type_field_name,
+                patched=is_patched_serializer(sub_serializer, direction)
+            )
+            sub_components.append((resource_type, typed_component.ref))
 
             if not resource_type:
                 warn(
@@ -32,3 +42,28 @@ def map_serializer(self, auto_schema, direction):
                 'mapping': {resource_type: ref['$ref'] for resource_type, ref in sub_components},
             }
         }
+
+    def build_typed_component(self, auto_schema, component, resource_type_field_name, patched):
+        if spectacular_settings.COMPONENT_SPLIT_REQUEST and component.name.endswith('Request'):
+            typed_component_name = component.name[:-len('Request')] + 'TypedRequest'
+        else:
+            typed_component_name = f'{component.name}Typed'
+
+        component_typed = ResolvedComponent(
+            name=typed_component_name,
+            type=ResolvedComponent.SCHEMA,
+            object=component.object,
+            schema={
+                'allOf': [
+                    component.ref,
+                    build_object_type(
+                        properties={
+                            resource_type_field_name: build_basic_type(OpenApiTypes.STR)
+                        },
+                        required=None if patched else [resource_type_field_name]
+                    )
+                ]
+            }
+        )
+        auto_schema.registry.register_on_missing(component_typed)
+        return component_typed
diff --git a/tests/contrib/test_rest_polymorphic.py b/tests/contrib/test_rest_polymorphic.py
@@ -101,22 +101,22 @@ def test_rest_polymorphic_split_request_with_ro_serializer(no_warnings):
     assert 'NomadicPersonRequest' not in components  # All fields were read-only.
     assert 'PatchedNomadicPersonRequest' not in components  # All fields were read-only.
     assert components['Person']['oneOf'] == [
-        {'$ref': '#/components/schemas/LegalPerson'},
-        {'$ref': '#/components/schemas/NaturalPerson'},
-        {'$ref': '#/components/schemas/NomadicPerson'}
+        {'$ref': '#/components/schemas/LegalPersonTyped'},
+        {'$ref': '#/components/schemas/NaturalPersonTyped'},
+        {'$ref': '#/components/schemas/NomadicPersonTyped'}
     ]
     assert components['Person']['discriminator']['mapping'] == {
-        'legal': '#/components/schemas/LegalPerson',
-        'natural': '#/components/schemas/NaturalPerson',
-        'nomadic': '#/components/schemas/NomadicPerson'
+        'legal': '#/components/schemas/LegalPersonTyped',
+        'natural': '#/components/schemas/NaturalPersonTyped',
+        'nomadic': '#/components/schemas/NomadicPersonTyped'
     }
     assert components['PersonRequest']['oneOf'] == [
-        {'$ref': '#/components/schemas/LegalPersonRequest'},
-        {'$ref': '#/components/schemas/NaturalPersonRequest'},
+        {'$ref': '#/components/schemas/LegalPersonTypedRequest'},
+        {'$ref': '#/components/schemas/NaturalPersonTypedRequest'},
     ]
     assert components['PersonRequest']['discriminator']['mapping'] == {
-        'legal': '#/components/schemas/LegalPersonRequest',
-        'natural': '#/components/schemas/NaturalPersonRequest',
+        'legal': '#/components/schemas/LegalPersonTypedRequest',
+        'natural': '#/components/schemas/NaturalPersonTypedRequest',
     }
 
 

diff --git a/tests/contrib/test_rest_polymorphic.yml b/tests/contrib/test_rest_polymorphic.yml
@@ -178,6 +178,15 @@ components:
       - board
       - company_name
       - id
+    LegalPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/LegalPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
+        required:
+        - resourcetype
     NaturalPerson:
       type: object
       properties:
@@ -202,6 +211,15 @@ components:
       - id
       - last_name
       - supervisor_id
+    NaturalPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/NaturalPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
+        required:
+        - resourcetype
     NomadicPerson:
       type: object
       properties:
@@ -215,6 +233,15 @@ components:
       required:
       - address
       - id
+    NomadicPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/NomadicPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
+        required:
+        - resourcetype
     PatchedLegalPerson:
       type: object
       properties:
@@ -232,6 +259,13 @@ components:
           items:
             $ref: '#/components/schemas/Person'
           readOnly: true
+    PatchedLegalPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/PatchedLegalPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
     PatchedNaturalPerson:
       type: object
       properties:
@@ -250,6 +284,13 @@ components:
         supervisor_id:
           type: integer
           nullable: true
+    PatchedNaturalPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/PatchedNaturalPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
     PatchedNomadicPerson:
       type: object
       properties:
@@ -260,28 +301,35 @@ components:
           type: string
           readOnly: true
           maxLength: 30
+    PatchedNomadicPersonTyped:
+      allOf:
+      - $ref: '#/components/schemas/PatchedNomadicPerson'
+      - type: object
+        properties:
+          resourcetype:
+            type: string
     PatchedPerson:
       oneOf:
-      - $ref: '#/components/schemas/PatchedLegalPerson'
-      - $ref: '#/components/schemas/PatchedNaturalPerson'
-      - $ref: '#/components/schemas/PatchedNomadicPerson'
+      - $ref: '#/components/schemas/PatchedLegalPersonTyped'
+      - $ref: '#/components/schemas/PatchedNaturalPersonTyped'
+      - $ref: '#/components/schemas/PatchedNomadicPersonTyped'
       discriminator:
         propertyName: resourcetype
         mapping:
-          legal: '#/components/schemas/PatchedLegalPerson'
-          natural: '#/components/schemas/PatchedNaturalPerson'
-          nomadic: '#/components/schemas/PatchedNomadicPerson'
+          legal: '#/components/schemas/PatchedLegalPersonTyped'
+          natural: '#/components/schemas/PatchedNaturalPersonTyped'
+          nomadic: '#/components/schemas/PatchedNomadicPersonTyped'
     Person:
       oneOf:
-      - $ref: '#/components/schemas/LegalPerson'
-      - $ref: '#/components/schemas/NaturalPerson'
-      - $ref: '#/components/schemas/NomadicPerson'
+      - $ref: '#/components/schemas/LegalPersonTyped'
+      - $ref: '#/components/schemas/NaturalPersonTyped'
+      - $ref: '#/components/schemas/NomadicPersonTyped'
       discriminator:
         propertyName: resourcetype
         mapping:
-          legal: '#/components/schemas/LegalPerson'
-          natural: '#/components/schemas/NaturalPerson'
-          nomadic: '#/components/schemas/NomadicPerson'
+          legal: '#/components/schemas/LegalPersonTyped'
+          natural: '#/components/schemas/NaturalPersonTyped'
+          nomadic: '#/components/schemas/NomadicPersonTyped'
   securitySchemes:
     basicAuth:
       type: http