From d628df6ab6c57b34acfb412e135d0d095636c539 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Sun, 20 Nov 2011 16:54:40 +0100
Subject: [PATCH] Store session after callbacks.  This fixes #351

---
 CHANGES                  |  3 +++
 flask/app.py             |  4 ++--
 flask/testsuite/basic.py | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/CHANGES b/CHANGES
index e30979d8..61723dca 100644
--- a/CHANGES
+++ b/CHANGES
@@ -20,6 +20,9 @@ Relase date to be decided, codename to be chosen.
   returned to the WSGI server but has the advantage that the garbage
   collector is not needed on CPython to tear down the request unless
   the user created circular dependencies themselves.
+- Session is now stored after callbacks so that if the session payload
+  is stored in the session you can still modify it in an after
+  request callback.
 
 Version 0.8.1
 -------------
diff --git a/flask/app.py b/flask/app.py
index ebf4e6a6..42ffea4d 100644
--- a/flask/app.py
+++ b/flask/app.py
@@ -1403,8 +1403,6 @@ class Flask(_PackageBoundObject):
         """
         ctx = _request_ctx_stack.top
         bp = ctx.request.blueprint
-        if not self.session_interface.is_null_session(ctx.session):
-            self.save_session(ctx.session, response)
         funcs = ()
         if bp is not None and bp in self.after_request_funcs:
             funcs = reversed(self.after_request_funcs[bp])
@@ -1412,6 +1410,8 @@ class Flask(_PackageBoundObject):
             funcs = chain(funcs, reversed(self.after_request_funcs[None]))
         for handler in funcs:
             response = handler(response)
+        if not self.session_interface.is_null_session(ctx.session):
+            self.save_session(ctx.session, response)
         return response
 
     def do_teardown_request(self):
diff --git a/flask/testsuite/basic.py b/flask/testsuite/basic.py
index 1733f0a3..a11f7806 100644
--- a/flask/testsuite/basic.py
+++ b/flask/testsuite/basic.py
@@ -279,6 +279,23 @@ class BasicFunctionalityTestCase(FlaskTestCase):
         match = re.search(r'\bexpires=([^;]+)', rv.headers['set-cookie'])
         self.assert_(match is None)
 
+    def test_session_stored_last(self):
+        app = flask.Flask(__name__)
+        app.secret_key = 'development-key'
+        app.testing = True
+
+        @app.after_request
+        def modify_session(response):
+            flask.session['foo'] = 42
+            return response
+        @app.route('/')
+        def dump_session_contents():
+            return repr(flask.session.get('foo'))
+
+        c = app.test_client()
+        self.assert_equal(c.get('/').data, 'None')
+        self.assert_equal(c.get('/').data, '42')
+
     def test_flashes(self):
         app = flask.Flask(__name__)
         app.secret_key = 'testkey'