diff --git a/flask_rest_api/etag.py b/flask_rest_api/etag.py index 49ff47c4..52f9faa1 100644 --- a/flask_rest_api/etag.py +++ b/flask_rest_api/etag.py @@ -85,9 +85,12 @@ def wrapper(*args, **kwargs): # Pass data to use as ETag data if set_etag was not called # If etag_schema is provided, pass raw result rather than # dump, as the dump needs to be done using etag_schema - etag_data = get_appcontext()[ - 'result_dump' if etag_schema is None else 'result_raw' - ] + # If 'result_dump'/'result_raw' is not in appcontext, + # the Etag must have been set manually. Just pass None. + etag_data = get_appcontext().get( + 'result_dump' if etag_schema is None else 'result_raw', + None + ) self._set_etag_in_response(resp, etag_data, etag_schema) return resp diff --git a/flask_rest_api/response.py b/flask_rest_api/response.py index 9d3a8767..7ad29282 100644 --- a/flask_rest_api/response.py +++ b/flask_rest_api/response.py @@ -2,9 +2,12 @@ from functools import wraps -from flask import jsonify +from flask import jsonify, Response -from .utils import deepupdate, get_appcontext, unpack_tuple_response +from .utils import ( + deepupdate, get_appcontext, + unpack_tuple_response, set_status_and_headers_in_response +) from .compat import MARSHMALLOW_VERSION_MAJOR @@ -16,7 +19,8 @@ def response(self, schema=None, *, code=200, description=''): :param schema: :class:`Schema ` class or instance. If not None, will be used to serialize response data. - :param int code: HTTP status code (default: 200). + :param int code: HTTP status code (default: 200). Used if none is + returned from the view function. :param str descripton: Description of the response. See :doc:`Response `. @@ -36,10 +40,18 @@ def decorator(func): @wraps(func) def wrapper(*args, **kwargs): + appcontext = get_appcontext() + # Execute decorated function result_raw, status, headers = unpack_tuple_response( func(*args, **kwargs)) + # If return value is a flask Response, return it + if isinstance(result_raw, Response): + set_status_and_headers_in_response( + result_raw, status, headers) + return result_raw + # Dump result with schema if specified if schema is None: result_dump = result_raw @@ -49,20 +61,14 @@ def wrapper(*args, **kwargs): result_dump = result_dump[0] # Store result in appcontext (may be used for ETag computation) - get_appcontext()['result_raw'] = result_raw - get_appcontext()['result_dump'] = result_dump + appcontext['result_raw'] = result_raw + appcontext['result_dump'] = result_dump # Build response resp = jsonify(self._prepare_response_content(result_dump)) - resp.headers.extend(get_appcontext()['headers']) - if headers: - resp.headers.extend(headers) - if status is not None: - if isinstance(status, int): - resp.status_code = status - else: - resp.status = status - else: + resp.headers.extend(appcontext['headers']) + set_status_and_headers_in_response(resp, status, headers) + if status is None: resp.status_code = code return resp diff --git a/flask_rest_api/utils.py b/flask_rest_api/utils.py index 14a3f0b6..3c44caf1 100644 --- a/flask_rest_api/utils.py +++ b/flask_rest_api/utils.py @@ -93,3 +93,14 @@ def unpack_tuple_response(rv): ) return rv, status, headers + + +def set_status_and_headers_in_response(response, status, headers): + """Set status and headers in flask Reponse object""" + if headers: + response.headers.extend(headers) + if status is not None: + if isinstance(status, int): + response.status_code = status + else: + response.status = status diff --git a/tests/test_blueprint.py b/tests/test_blueprint.py index 5460028b..95e844e5 100644 --- a/tests/test_blueprint.py +++ b/tests/test_blueprint.py @@ -604,3 +604,22 @@ def func_response_wrong_tuple(): assert response.headers['X-header'] == 'test' response = client.get('/test/response_wrong_tuple') assert response.status_code == 500 + + def test_blueprint_response_response_object(self, app, schemas): + api = Api(app) + blp = Blueprint('test', __name__, url_prefix='/test') + client = app.test_client() + + @blp.route('/response') + # Schema is ignored when response object is returned + @blp.response(schemas.DocSchema, code=200) + def func_response(): + return jsonify({}), 201, {'X-header': 'test'} + + api.register_blueprint(blp) + + response = client.get('/test/response') + assert response.status_code == 201 + assert response.status == '201 CREATED' + assert response.json == {} + assert response.headers['X-header'] == 'test' diff --git a/tests/test_etag.py b/tests/test_etag.py index 0fe2acd4..62bff595 100644 --- a/tests/test_etag.py +++ b/tests/test_etag.py @@ -7,7 +7,7 @@ import pytest -from flask import Response +from flask import jsonify, Response from flask.views import MethodView from flask_rest_api import Api, Blueprint, abort @@ -358,6 +358,26 @@ def test_etag_set_etag_in_response(self, app, schemas, paginate): blp._set_etag_in_response(resp, item, etag_schema) assert resp.get_etag() == (etag_with_schema, False) + def test_etag_response_object(self, app): + api = Api(app) + blp = Blueprint('test', __name__, url_prefix='/test') + client = app.test_client() + + @blp.route('/') + @blp.etag + @blp.response() + def func_response_etag(): + # When the view function returns a Response object, + # the ETag must be specified manually + blp.set_etag('test') + return jsonify({}) + + api.register_blueprint(blp) + + response = client.get('/test/') + assert response.json == {} + assert response.get_etag() == (blp._generate_etag('test'), False) + def test_etag_operations_etag_enabled(self, app_with_etag): client = app_with_etag.test_client()