Merge pull request #184 from c-kruse/feature/marshal_with_returns_response

Return Response Object from View Functions

Author: sloria

CHANGELOG.rst

@@ -1,6 +1,16 @@
 Changelog
 ---------
 
+0.8.8 (unreleased)
+******************
+
+Bug fixes:
+
+* Fix behavior when view returns ``(obj, status_code, headers)``
+  (regression in 0.8.7) (:issue:`181`).
+  Thanks :user:`decaz` for reporting and thanks :user:`c-kruse`
+  for the PR.
+
 0.8.7 (2020-03-10)
 ******************
 

flask_apispec/wrapper.py

@@ -34,9 +34,10 @@ def __call__(self, *args, **kwargs):
         response = self.call_view(*args, **kwargs)
         if isinstance(response, werkzeug.Response):
             return response
-        unpacked = unpack(response)
-        status_code = unpacked[1] or http.OK
-        return self.marshal_result(unpacked, status_code)
+        rv, status_code, headers = unpack(response)
+        mv = self.marshal_result(rv, status_code)
+        response = packed(mv, status_code, headers)
+        return flask.current_app.make_response(response)
 
     def call_view(self, *args, **kwargs):
         config = flask.current_app.config
@@ -55,34 +56,48 @@ def call_view(self, *args, **kwargs):
 
         return self.func(*args, **kwargs)
 
-    def marshal_result(self, unpacked, status_code):
+    def marshal_result(self, result, status_code):
         config = flask.current_app.config
         format_response = config.get('APISPEC_FORMAT_RESPONSE', flask.jsonify) or identity
         annotation = utils.resolve_annotations(self.func, 'schemas', self.instance)
         schemas = utils.merge_recursive(annotation.options)
         schema = schemas.get(status_code, schemas.get('default'))
         if schema and annotation.apply is not False:
             schema = utils.resolve_schema(schema['schema'], request=flask.request)
-            dumped = schema.dump(unpacked[0])
+            dumped = schema.dump(result)
             output = dumped.data if MARSHMALLOW_VERSION_INFO[0] < 3 else dumped
         else:
-            output = unpacked[0]
+            output = result
 
-        response_object = format_response(output)  # type: Response
-        response_object.status_code = unpacked[1] or 200
-        return format_output((response_object,))
+        return format_response(output) # type: Response
 
 
 def identity(value):
     return value
 
 
 def unpack(resp):
-    resp = resp if isinstance(resp, tuple) else (resp,)
-    return resp + (None,) * (3 - len(resp))
-
-
-def format_output(values):
-    while values[-1] is None:
-        values = values[:-1]
-    return values if len(values) > 1 else values[0]
+    data = headers = None
+    status_code = 200
+    if isinstance(resp, tuple):
+        data = resp[0]
+        len_resp = len(resp)
+        if len_resp == 3:
+            status_code, headers = resp[1:]
+        elif len_resp == 2:
+            if isinstance(resp[1], (werkzeug.datastructures.Headers, dict, tuple, list)):
+                headers = resp[1]
+            else:
+                status_code = resp[1]
+    else:
+        data = resp
+    return data, status_code, headers
+
+
+def packed(data, status_code, headers):
+    resp = (data, )
+    if status_code:
+        resp += (status_code, )
+    if headers:
+        resp += (headers, )
+    return resp

tests/test_views.py

@@ -19,6 +19,37 @@ def view(**kwargs):
         res = client.get('/', {'name': 'freddie'})
         assert res.json == {'name': 'freddie'}
 
+    def test_view_returning_tuple(self, app, client):
+        @app.route('/all')
+        @use_kwargs({'name': fields.Str()})
+        def all(**kwargs):
+            return kwargs, 202, {'x-msg': 'test'}
+
+        @app.route('/headers')
+        @use_kwargs({'name': fields.Str()})
+        def view_headers(**kwargs):
+            return kwargs, {'x-msg': 'test'}
+
+        @app.route('/code')
+        @use_kwargs({'name': fields.Str()})
+        def view_code(**kwargs):
+            return kwargs, 202
+
+        res_all = client.get('/all', {'name': 'freddie'})
+        assert res_all.json == {'name': 'freddie'}
+        assert res_all.status_code == 202
+        assert res_all.headers.get('x-msg') == 'test'
+
+        res_headers = client.get('/headers', {'name': 'freddie'})
+        assert res_headers.json == {'name': 'freddie'}
+        assert res_headers.status_code == 200
+        assert res_headers.headers.get('x-msg') == 'test'
+
+        res_code = client.get('/code', {'name': 'freddie'})
+        assert res_code.json == {'name': 'freddie'}
+        assert res_code.status_code == 202
+        assert 'x-msg' not in res_code.headers
+
     def test_use_kwargs_schema(self, app, client):
         class ArgSchema(Schema):
             name = fields.Str()