Added flask.stream_with_context

CHANGES

@@ -71,6 +71,8 @@ Relase date to be decided, codename to be chosen.
 - Added `required_methods` attribute to view functions to force-add methods
   on registration.
 - Added :func:`flask.after_this_request`.
+- Added :func:`flask.stream_with_context` and the ability to push contexts
+  multiple times without producing unexpected behavior.
 
 Version 0.8.1
 -------------

docs/api.rst

@@ -375,6 +375,11 @@ Extensions
 
    .. versionadded:: 0.8
 
+Stream Helpers
+--------------
+
+.. autofunction:: stream_with_context
+
 Useful Internals
 ----------------
 

docs/patterns/streaming.rst

@@ -59,3 +59,26 @@ The template is then evaluated as the stream is iterated over.  Since each
 time you do a yield the server will flush the content to the client you
 might want to buffer up a few items in the template which you can do with
 ``rv.enable_buffering(size)``.  ``5`` is a sane default.
+
+Streaming with Context
+----------------------
+
+.. versionadded:: 0.9
+
+Note that when you stream data, the request context is already gone the
+moment the function executes.  Flask 0.9 provides you with a helper that
+can keep the request context around during the execution of the
+generator::
+
+    from flask import stream_with_context, request, Response
+
+    @app.route('/stream')
+    def streamed_response():
+        def generate():
+            yield 'Hello '
+            yield request.args['name']
+            yield '!'
+        return Response(stream_with_context(generate()))
+
+Without the :func:`~flask.stream_with_context` function you would get a
+:class:`RuntimeError` at that point.

flask/__init__.py

@@ -22,7 +22,8 @@ from .app import Flask, Request, Response
 from .config import Config
 from .helpers import url_for, jsonify, json_available, flash, \
     send_file, send_from_directory, get_flashed_messages, \
-    get_template_attribute, make_response, safe_join
+    get_template_attribute, make_response, safe_join, \
+    stream_with_context
 from .globals import current_app, g, request, session, _request_ctx_stack, \
      _app_ctx_stack
 from .ctx import has_request_context, has_app_context, \

flask/ctx.py

@@ -22,14 +22,6 @@ class _RequestGlobals(object):
     pass
 
 
-def _push_app_if_necessary(app):
-    top = _app_ctx_stack.top
-    if top is None or top.app != app:
-        ctx = app.app_context()
-        ctx.push()
-        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
@@ -110,12 +102,19 @@ class AppContext(object):
         self.app = app
         self.url_adapter = app.create_url_adapter(None)
 
+        # Like request context, app contexts can be pushed multiple times
+        # but there a basic "refcount" is enough to track them.
+        self._refcnt = 0
+
     def push(self):
         """Binds the app context to the current context."""
+        self._refcnt += 1
         _app_ctx_stack.push(self)
 
     def pop(self, exc=None):
         """Pops the app context."""
+        self._refcnt -= 1
+        if self._refcnt <= 0:
             if exc is None:
                 exc = sys.exc_info()[1]
             self.app.do_teardown_appcontext(exc)
@@ -128,7 +127,7 @@ class AppContext(object):
         return self
 
     def __exit__(self, exc_type, exc_value, tb):
-        self.pop()
+        self.pop(exc_value)
 
 
 class RequestContext(object):
@@ -169,15 +168,16 @@ class RequestContext(object):
         self.flashes = None
         self.session = None
 
+        # Request contexts can be pushed multiple times and interleaved with
+        # other request contexts.  Now only if the last level is popped we
+        # get rid of them.  Additionally if an application context is missing
+        # one is created implicitly so for each level we add this information
+        self._implicit_app_ctx_stack = []
+
         # indicator if the context was preserved.  Next time another context
         # is pushed the preserved context is popped.
         self.preserved = False
 
-        # Indicates if pushing this request context also triggered the pushing
-        # of an application context.  If it implicitly pushed an application
-        # 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 be called before the regular "after_request"
         # functions.
@@ -222,7 +222,13 @@ class RequestContext(object):
 
         # Before we push the request context we have to ensure that there
         # is an application context.
-        self._pushed_application_context = _push_app_if_necessary(self.app)
+        app_ctx = _app_ctx_stack.top
+        if app_ctx is None or app_ctx.app != self.app:
+            app_ctx = self.app.app_context()
+            app_ctx.push()
+            self._implicit_app_ctx_stack.append(app_ctx)
+        else:
+            self._implicit_app_ctx_stack.append(None)
 
         _request_ctx_stack.push(self)
 
@@ -241,22 +247,28 @@ class RequestContext(object):
         .. versionchanged:: 0.9
            Added the `exc` argument.
         """
+        app_ctx = self._implicit_app_ctx_stack.pop()
+
+        clear_request = False
+        if not self._implicit_app_ctx_stack:
             self.preserved = False
             if exc is None:
                 exc = sys.exc_info()[1]
             self.app.do_teardown_request(exc)
+            clear_request = True
+
         rv = _request_ctx_stack.pop()
         assert rv is self, 'Popped wrong request context.  (%r instead of %r)' \
             % (rv, self)
 
         # get rid of circular dependencies at the end of the request
         # so that we don't require the GC to be active.
+        if clear_request:
             rv.request.environ['werkzeug.request'] = None
 
         # Get rid of the app as well if necessary.
-        if self._pushed_application_context:
-            self._pushed_application_context.pop(exc)
-            self._pushed_application_context = None
+        if app_ctx is not None:
+            app_ctx.pop(exc)
 
     def __enter__(self):
         self.push()

flask/helpers.py

@@ -21,6 +21,7 @@ from zlib import adler32
 from threading import RLock
 from werkzeug.routing import BuildError
 from werkzeug.urls import url_quote
+from functools import update_wrapper
 
 # try to load the best simplejson implementation available.  If JSON
 # is not installed, we add a failing class.
@@ -92,6 +93,78 @@ def _endpoint_from_view_func(view_func):
     return view_func.__name__
 
 
+def stream_with_context(generator_or_function):
+    """Request contexts disappear when the response is started on the server.
+    This is done for efficiency reasons and to make it less likely to encounter
+    memory leaks with badly written WSGI middlewares.  The downside is that if
+    you are using streamed responses, the generator cannot access request bound
+    information any more.
+
+    This function however can help you keep the context around for longer::
+
+        from flask import stream_with_context, request, Response
+
+        @app.route('/stream')
+        def streamed_response():
+            @stream_with_context
+            def generate():
+                yield 'Hello '
+                yield request.args['name']
+                yield '!'
+            return Response(generate())
+
+    Alternatively it can also be used around a specific generator:
+
+        from flask import stream_with_context, request, Response
+
+        @app.route('/stream')
+        def streamed_response():
+            def generate():
+                yield 'Hello '
+                yield request.args['name']
+                yield '!'
+            return Response(stream_with_context(generate()))
+
+    .. versionadded:: 0.9
+    """
+    try:
+        gen = iter(generator_or_function)
+    except TypeError:
+        def decorator(*args, **kwargs):
+            gen = generator_or_function()
+            return stream_with_context(gen)
+        return update_wrapper(decorator, generator_or_function)
+
+    def generator():
+        ctx = _request_ctx_stack.top
+        if ctx is None:
+            raise RuntimeError('Attempted to stream with context but '
+                'there was no context in the first place to keep around.')
+        with ctx:
+            # Dummy sentinel.  Has to be inside the context block or we're
+            # not actually keeping the context around.
+            yield None
+
+            # The try/finally is here so that if someone passes a WSGI level
+            # iterator in we're still running the cleanup logic.  Generators
+            # don't need that because they are closed on their destruction
+            # automatically.
+            try:
+                for item in gen:
+                    yield item
+            finally:
+                if hasattr(gen, 'close'):
+                    gen.close()
+
+    # The trick is to start the generator.  Then the code execution runs until
+    # the first dummy None is yielded at which point the context was already
+    # pushed.  This item is discarded.  Then when the iteration continues the
+    # real generator is executed.
+    wrapped_g = generator()
+    wrapped_g.next()
+    return wrapped_g
+
+
 def jsonify(*args, **kwargs):
     """Creates a :class:`~flask.Response` with the JSON representation of
     the given arguments with an `application/json` mimetype.  The arguments

flask/testsuite/appctx.py

@@ -75,6 +75,26 @@ class AppContextTestCase(FlaskTestCase):
             self.assert_equal(
                 flask.render_template_string('{{ g.spam }}'), 'eggs')
 
+    def test_context_refcounts(self):
+        called = []
+        app = flask.Flask(__name__)
+        @app.teardown_request
+        def teardown_req(error=None):
+            called.append('request')
+        @app.teardown_appcontext
+        def teardown_app(error=None):
+            called.append('app')
+        @app.route('/')
+        def index():
+            with flask._app_ctx_stack.top:
+                with flask._request_ctx_stack.top:
+                    pass
+            self.assert_(flask._request_ctx_stack.request.environ
+                ['werkzeug.request'] is not None)
+        c = app.test_client()
+        c.get('/')
+        self.assertEqual(called, ['request', 'app'])
+
 
 def suite():
     suite = unittest.TestSuite()

flask/testsuite/helpers.py

@@ -397,6 +397,64 @@ class NoImportsTestCase(FlaskTestCase):
             self.fail('Flask(import_name) is importing import_name.')
 
 
+class StreamingTestCase(FlaskTestCase):
+
+    def test_streaming_with_context(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        @app.route('/')
+        def index():
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(flask.stream_with_context(generate()))
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+
+    def test_streaming_with_context_as_decorator(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        @app.route('/')
+        def index():
+            @flask.stream_with_context
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(generate())
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+
+    def test_streaming_with_context_and_custom_close(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        called = []
+        class Wrapper(object):
+            def __init__(self, gen):
+                self._gen = gen
+            def __iter__(self):
+                return self
+            def close(self):
+                called.append(42)
+            def next(self):
+                return self._gen.next()
+        @app.route('/')
+        def index():
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(flask.stream_with_context(
+                Wrapper(generate())))
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+        self.assertEqual(called, [42])
+
+
 def suite():
     suite = unittest.TestSuite()
     if flask.json_available:
@@ -404,4 +462,5 @@ def suite():
     suite.addTest(unittest.makeSuite(SendfileTestCase))
     suite.addTest(unittest.makeSuite(LoggingTestCase))
     suite.addTest(unittest.makeSuite(NoImportsTestCase))
+    suite.addTest(unittest.makeSuite(StreamingTestCase))
     return suite