From 086348e2f2874fb701048d5e1390cfe674de1f70 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Tue, 8 May 2012 13:14:32 +0100
Subject: [PATCH] Added after_this_request decorator.

---
 CHANGES                  |  1 +
 docs/api.rst             |  2 ++
 flask/__init__.py        |  3 ++-
 flask/app.py             |  2 +-
 flask/ctx.py             | 30 ++++++++++++++++++++++++++++++
 flask/testsuite/basic.py | 15 +++++++++++++++
 6 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/CHANGES b/CHANGES
index 2c3f904a..28a0790c 100644
--- a/CHANGES
+++ b/CHANGES
@@ -70,6 +70,7 @@ Relase date to be decided, codename to be chosen.
   be used on creation of the :data:`~flask.g` instance of each request.
 - Added `required_methods` attribute to view functions to force-add methods
   on registration.
+- Added :func:`flask.after_this_request`.
 
 Version 0.8.1
 -------------
diff --git a/docs/api.rst b/docs/api.rst
index c97a6928..dcb54baf 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -289,6 +289,8 @@ Useful Functions and Classes
 
 .. autofunction:: make_response
 
+.. autofunction:: after_this_request
+
 .. autofunction:: send_file
 
 .. autofunction:: send_from_directory
diff --git a/flask/__init__.py b/flask/__init__.py
index b91f9395..de84bb69 100644
--- a/flask/__init__.py
+++ b/flask/__init__.py
@@ -25,7 +25,8 @@ from .helpers import url_for, jsonify, json_available, flash, \
     get_template_attribute, make_response, safe_join
 from .globals import current_app, g, request, session, _request_ctx_stack, \
      _app_ctx_stack
-from .ctx import has_request_context, has_app_context
+from .ctx import has_request_context, has_app_context, \
+     after_this_request
 from .module import Module
 from .blueprints import Blueprint
 from .templating import render_template, render_template_string
diff --git a/flask/app.py b/flask/app.py
index a91b0f60..9254c398 100644
--- a/flask/app.py
+++ b/flask/app.py
@@ -1555,7 +1555,7 @@ class Flask(_PackageBoundObject):
         """
         ctx = _request_ctx_stack.top
         bp = ctx.request.blueprint
-        funcs = ()
+        funcs = ctx._after_request_functions
         if bp is not None and bp in self.after_request_funcs:
             funcs = reversed(self.after_request_funcs[bp])
         if None in self.after_request_funcs:
diff --git a/flask/ctx.py b/flask/ctx.py
index cf197d05..52eddcb2 100644
--- a/flask/ctx.py
+++ b/flask/ctx.py
@@ -11,6 +11,7 @@
 
 import sys
 
+from functools import partial
 from werkzeug.exceptions import HTTPException
 
 from .globals import _request_ctx_stack, _app_ctx_stack
@@ -30,6 +31,31 @@ def _push_app_if_necessary(app):
         return ctx
 
 
+def after_this_request(f):
+    """Executes a function after this request.  This is useful to modify
+    response objects.  The function is passed the response object and has
+    to return the same or a new one.
+
+    Example::
+
+        @app.route('/')
+        def index():
+            @after_this_request
+            def add_header():
+                response.headers['X-Foo'] = 'Parachute'
+                return response
+            return 'Hello World!'
+
+    This is more useful if a function other than the view function wants to
+    modify a response.  For instance think of a decorator that wants to add
+    some headers without converting the return value into a response object.
+
+    .. versionadded:: 0.9
+    """
+    _request_ctx_stack.top._after_request_functions.append(f)
+    return f
+
+
 def has_request_context():
     """If you have code that wants to test if a request context is there or
     not this function can be used.  For instance, you may want to take advantage
@@ -153,6 +179,10 @@ class RequestContext(object):
         # context, it will be stored there
         self._pushed_application_context = None
 
+        # Functions that should be executed after the request on the response
+        # object.  These will even be called in case of an error.
+        self._after_request_functions = []
+
         self.match_request()
 
         # XXX: Support for deprecated functionality.  This is going away with
diff --git a/flask/testsuite/basic.py b/flask/testsuite/basic.py
index 55b66f78..ba6c2705 100644
--- a/flask/testsuite/basic.py
+++ b/flask/testsuite/basic.py
@@ -411,6 +411,21 @@ class BasicFunctionalityTestCase(FlaskTestCase):
         self.assert_('after' in evts)
         self.assert_equal(rv, 'request|after')
 
+    def test_after_request_processing(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        @app.route('/')
+        def index():
+            @flask.after_this_request
+            def foo(response):
+                response.headers['X-Foo'] = 'a header'
+                return response
+            return 'Test'
+        c = app.test_client()
+        resp = c.get('/')
+        self.assertEqual(resp.status_code, 200)
+        self.assertEqual(resp.headers['X-Foo'], 'a header')
+
     def test_teardown_request_handler(self):
         called = []
         app = flask.Flask(__name__)