Merge branch 'pr301'

tfranzel · Feb 17, 2021 · d76c310 · d76c310
2 parents f219b6e + 5e2b262
commit d76c310
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 2 deletions.
diff --git a/drf_spectacular/openapi.py b/drf_spectacular/openapi.py
@@ -124,6 +124,9 @@ def _process_override_parameters(self):
         result = {}
         for parameter in self.get_override_parameters():
             if isinstance(parameter, OpenApiParameter):
+                if parameter.response:
+                    continue
+
                 if is_basic_type(parameter.type):
                     schema = build_basic_type(parameter.type)
                 elif is_serializer(parameter.type):
@@ -1003,7 +1006,7 @@ def _get_response_for_code(self, serializer, status_code, media_types=None):
         if not media_types:
             media_types = self.map_renderers('media_type')
 
-        return {
+        response = {
             'content': {
                 media_type: build_media_type_object(
                     schema,
@@ -1016,6 +1019,52 @@ def _get_response_for_code(self, serializer, status_code, media_types=None):
             # https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.3.md#responseObject
             'description': ''
         }
+        headers = self._get_response_headers_for_code(status_code)
+        if headers:
+            response['headers'] = headers
+        return response
+
+    def _get_response_headers_for_code(self, status_code) -> dict:
+        result = {}
+        for parameter in self.get_override_parameters():
+            if not isinstance(parameter, OpenApiParameter):
+                continue
+            if not parameter.response:
+                continue
+            if (
+                isinstance(parameter.response, list)
+                and status_code not in [str(code) for code in parameter.response]
+            ):
+                continue
+
+            if is_basic_type(parameter.type):
+                schema = build_basic_type(parameter.type)
+            elif is_serializer(parameter.type):
+                schema = self.resolve_serializer(parameter.type, 'response').ref
+            else:
+                schema = parameter.type
+
+            if parameter.location not in [OpenApiParameter.HEADER, OpenApiParameter.COOKIE]:
+                warn(f'incompatible location type ignored for response parameter {parameter.name}')
+
+            parameter_type = build_parameter_type(
+                name=parameter.name,
+                schema=schema,
+                location=parameter.location,
+                required=parameter.required,
+                description=parameter.description,
+                enum=parameter.enum,
+                deprecated=parameter.deprecated,
+                style=parameter.style,
+                explode=parameter.explode,
+                default=parameter.default,
+                examples=build_examples_list(parameter.examples),
+            )
+            del parameter_type['name']
+            del parameter_type['in']
+            result[parameter.name] = parameter_type
+
+        return result
 
     def _get_serializer_name(self, serializer, direction):
         serializer_extension = OpenApiSerializerExtension.get_match(serializer)

diff --git a/drf_spectacular/utils.py b/drf_spectacular/utils.py
@@ -115,6 +115,7 @@ def __init__(
             default=None,
             examples: Optional[List[OpenApiExample]] = None,
             exclude=False,
+            response: Union[bool, List[Union[int, str]]] = False,
     ):
         self.name = name
         self.type = type
@@ -128,6 +129,7 @@ def __init__(
         self.default = default
         self.examples = examples or []
         self.exclude = exclude
+        self.response = response
 
 
 def extend_schema(

diff --git a/tests/test_extend_schema.py b/tests/test_extend_schema.py
@@ -98,7 +98,20 @@ class DoesItAllViewset(viewsets.GenericViewSet):
                     enum=[True, False],
                     default=False,
                     description='creation will be in the sandbox',
-                )
+                ),
+                OpenApiParameter(
+                    name='X-Api-Version',
+                    type=str,
+                    location=OpenApiParameter.HEADER,
+                    response=True,
+                ),
+                OpenApiParameter(
+                    name='Location',
+                    type=OpenApiTypes.URI,
+                    location=OpenApiParameter.HEADER,
+                    description='URL of the created resource',
+                    response=[201],
+                ),
             ],
             description='this weird endpoint needs some explaining',
             summary='short summary',

diff --git a/tests/test_extend_schema.yml b/tests/test_extend_schema.yml
@@ -52,18 +52,35 @@ paths:
                 items:
                   $ref: '#/components/schemas/Beta'
           description: ''
+          headers:
+            X-Api-Version:
+              schema:
+                type: string
         '201':
           content:
             application/json:
               schema:
                 $ref: '#/components/schemas/Gamma'
           description: ''
+          headers:
+            X-Api-Version:
+              schema:
+                type: string
+            Location:
+              schema:
+                type: string
+                format: uri
+              description: URL of the created resource
         '500':
           content:
             application/json:
               schema:
                 $ref: '#/components/schemas/ErrorDetail'
           description: ''
+          headers:
+            X-Api-Version:
+              schema:
+                type: string
   /doesitall/{id}/subscribe/:
     post:
       operationId: doesitall_subscribe_create