Skip to content
This repository
Browse code

Added flask.stream_with_context

  • Loading branch information...
commit d5218997d927be869dd55ef04542e1bbc1e69653 1 parent 2e816f5
Armin Ronacher authored June 27, 2012
2  CHANGES
@@ -71,6 +71,8 @@ Relase date to be decided, codename to be chosen.
71 71
 - Added `required_methods` attribute to view functions to force-add methods
72 72
   on registration.
73 73
 - Added :func:`flask.after_this_request`.
  74
+- Added :func:`flask.stream_with_context` and the ability to push contexts
  75
+  multiple times without producing unexpected behavior.
74 76
 
75 77
 Version 0.8.1
76 78
 -------------
5  docs/api.rst
Source Rendered
@@ -375,6 +375,11 @@ Extensions
375 375
 
376 376
    .. versionadded:: 0.8
377 377
 
  378
+Stream Helpers
  379
+--------------
  380
+
  381
+.. autofunction:: stream_with_context
  382
+
378 383
 Useful Internals
379 384
 ----------------
380 385
 
23  docs/patterns/streaming.rst
Source Rendered
@@ -59,3 +59,26 @@ The template is then evaluated as the stream is iterated over.  Since each
59 59
 time you do a yield the server will flush the content to the client you
60 60
 might want to buffer up a few items in the template which you can do with
61 61
 ``rv.enable_buffering(size)``.  ``5`` is a sane default.
  62
+
  63
+Streaming with Context
  64
+----------------------
  65
+
  66
+.. versionadded:: 0.9
  67
+
  68
+Note that when you stream data, the request context is already gone the
  69
+moment the function executes.  Flask 0.9 provides you with a helper that
  70
+can keep the request context around during the execution of the
  71
+generator::
  72
+
  73
+    from flask import stream_with_context, request, Response
  74
+
  75
+    @app.route('/stream')
  76
+    def streamed_response():
  77
+        def generate():
  78
+            yield 'Hello '
  79
+            yield request.args['name']
  80
+            yield '!'
  81
+        return Response(stream_with_context(generate()))
  82
+
  83
+Without the :func:`~flask.stream_with_context` function you would get a
  84
+:class:`RuntimeError` at that point.
3  flask/__init__.py
@@ -22,7 +22,8 @@
22 22
 from .config import Config
23 23
 from .helpers import url_for, jsonify, json_available, flash, \
24 24
     send_file, send_from_directory, get_flashed_messages, \
25  
-    get_template_attribute, make_response, safe_join
  25
+    get_template_attribute, make_response, safe_join, \
  26
+    stream_with_context
26 27
 from .globals import current_app, g, request, session, _request_ctx_stack, \
27 28
      _app_ctx_stack
28 29
 from .ctx import has_request_context, has_app_context, \
64  flask/ctx.py
@@ -22,14 +22,6 @@ class _RequestGlobals(object):
22 22
     pass
23 23
 
24 24
 
25  
-def _push_app_if_necessary(app):
26  
-    top = _app_ctx_stack.top
27  
-    if top is None or top.app != app:
28  
-        ctx = app.app_context()
29  
-        ctx.push()
30  
-        return ctx
31  
-
32  
-
33 25
 def after_this_request(f):
34 26
     """Executes a function after this request.  This is useful to modify
35 27
     response objects.  The function is passed the response object and has
@@ -110,15 +102,22 @@ def __init__(self, app):
110 102
         self.app = app
111 103
         self.url_adapter = app.create_url_adapter(None)
112 104
 
  105
+        # Like request context, app contexts can be pushed multiple times
  106
+        # but there a basic "refcount" is enough to track them.
  107
+        self._refcnt = 0
  108
+
113 109
     def push(self):
114 110
         """Binds the app context to the current context."""
  111
+        self._refcnt += 1
115 112
         _app_ctx_stack.push(self)
116 113
 
117 114
     def pop(self, exc=None):
118 115
         """Pops the app context."""
119  
-        if exc is None:
120  
-            exc = sys.exc_info()[1]
121  
-        self.app.do_teardown_appcontext(exc)
  116
+        self._refcnt -= 1
  117
+        if self._refcnt <= 0:
  118
+            if exc is None:
  119
+                exc = sys.exc_info()[1]
  120
+            self.app.do_teardown_appcontext(exc)
122 121
         rv = _app_ctx_stack.pop()
123 122
         assert rv is self, 'Popped wrong app context.  (%r instead of %r)' \
124 123
             % (rv, self)
@@ -128,7 +127,7 @@ def __enter__(self):
128 127
         return self
129 128
 
130 129
     def __exit__(self, exc_type, exc_value, tb):
131  
-        self.pop()
  130
+        self.pop(exc_value)
132 131
 
133 132
 
134 133
 class RequestContext(object):
@@ -169,15 +168,16 @@ def __init__(self, app, environ):
169 168
         self.flashes = None
170 169
         self.session = None
171 170
 
  171
+        # Request contexts can be pushed multiple times and interleaved with
  172
+        # other request contexts.  Now only if the last level is popped we
  173
+        # get rid of them.  Additionally if an application context is missing
  174
+        # one is created implicitly so for each level we add this information
  175
+        self._implicit_app_ctx_stack = []
  176
+
172 177
         # indicator if the context was preserved.  Next time another context
173 178
         # is pushed the preserved context is popped.
174 179
         self.preserved = False
175 180
 
176  
-        # Indicates if pushing this request context also triggered the pushing
177  
-        # of an application context.  If it implicitly pushed an application
178  
-        # context, it will be stored there
179  
-        self._pushed_application_context = None
180  
-
181 181
         # Functions that should be executed after the request on the response
182 182
         # object.  These will be called before the regular "after_request"
183 183
         # functions.
@@ -222,7 +222,13 @@ def push(self):
222 222
 
223 223
         # Before we push the request context we have to ensure that there
224 224
         # is an application context.
225  
-        self._pushed_application_context = _push_app_if_necessary(self.app)
  225
+        app_ctx = _app_ctx_stack.top
  226
+        if app_ctx is None or app_ctx.app != self.app:
  227
+            app_ctx = self.app.app_context()
  228
+            app_ctx.push()
  229
+            self._implicit_app_ctx_stack.append(app_ctx)
  230
+        else:
  231
+            self._implicit_app_ctx_stack.append(None)
226 232
 
227 233
         _request_ctx_stack.push(self)
228 234
 
@@ -241,22 +247,28 @@ def pop(self, exc=None):
241 247
         .. versionchanged:: 0.9
242 248
            Added the `exc` argument.
243 249
         """
244  
-        self.preserved = False
245  
-        if exc is None:
246  
-            exc = sys.exc_info()[1]
247  
-        self.app.do_teardown_request(exc)
  250
+        app_ctx = self._implicit_app_ctx_stack.pop()
  251
+
  252
+        clear_request = False
  253
+        if not self._implicit_app_ctx_stack:
  254
+            self.preserved = False
  255
+            if exc is None:
  256
+                exc = sys.exc_info()[1]
  257
+            self.app.do_teardown_request(exc)
  258
+            clear_request = True
  259
+
248 260
         rv = _request_ctx_stack.pop()
249 261
         assert rv is self, 'Popped wrong request context.  (%r instead of %r)' \
250 262
             % (rv, self)
251 263
 
252 264
         # get rid of circular dependencies at the end of the request
253 265
         # so that we don't require the GC to be active.
254  
-        rv.request.environ['werkzeug.request'] = None
  266
+        if clear_request:
  267
+            rv.request.environ['werkzeug.request'] = None
255 268
 
256 269
         # Get rid of the app as well if necessary.
257  
-        if self._pushed_application_context:
258  
-            self._pushed_application_context.pop(exc)
259  
-            self._pushed_application_context = None
  270
+        if app_ctx is not None:
  271
+            app_ctx.pop(exc)
260 272
 
261 273
     def __enter__(self):
262 274
         self.push()
73  flask/helpers.py
@@ -21,6 +21,7 @@
21 21
 from threading import RLock
22 22
 from werkzeug.routing import BuildError
23 23
 from werkzeug.urls import url_quote
  24
+from functools import update_wrapper
24 25
 
25 26
 # try to load the best simplejson implementation available.  If JSON
26 27
 # is not installed, we add a failing class.
@@ -92,6 +93,78 @@ def _endpoint_from_view_func(view_func):
92 93
     return view_func.__name__
93 94
 
94 95
 
  96
+def stream_with_context(generator_or_function):
  97
+    """Request contexts disappear when the response is started on the server.
  98
+    This is done for efficiency reasons and to make it less likely to encounter
  99
+    memory leaks with badly written WSGI middlewares.  The downside is that if
  100
+    you are using streamed responses, the generator cannot access request bound
  101
+    information any more.
  102
+
  103
+    This function however can help you keep the context around for longer::
  104
+
  105
+        from flask import stream_with_context, request, Response
  106
+
  107
+        @app.route('/stream')
  108
+        def streamed_response():
  109
+            @stream_with_context
  110
+            def generate():
  111
+                yield 'Hello '
  112
+                yield request.args['name']
  113
+                yield '!'
  114
+            return Response(generate())
  115
+
  116
+    Alternatively it can also be used around a specific generator:
  117
+
  118
+        from flask import stream_with_context, request, Response
  119
+
  120
+        @app.route('/stream')
  121
+        def streamed_response():
  122
+            def generate():
  123
+                yield 'Hello '
  124
+                yield request.args['name']
  125
+                yield '!'
  126
+            return Response(stream_with_context(generate()))
  127
+
  128
+    .. versionadded:: 0.9
  129
+    """
  130
+    try:
  131
+        gen = iter(generator_or_function)
  132
+    except TypeError:
  133
+        def decorator(*args, **kwargs):
  134
+            gen = generator_or_function()
  135
+            return stream_with_context(gen)
  136
+        return update_wrapper(decorator, generator_or_function)
  137
+
  138
+    def generator():
  139
+        ctx = _request_ctx_stack.top
  140
+        if ctx is None:
  141
+            raise RuntimeError('Attempted to stream with context but '
  142
+                'there was no context in the first place to keep around.')
  143
+        with ctx:
  144
+            # Dummy sentinel.  Has to be inside the context block or we're
  145
+            # not actually keeping the context around.
  146
+            yield None
  147
+
  148
+            # The try/finally is here so that if someone passes a WSGI level
  149
+            # iterator in we're still running the cleanup logic.  Generators
  150
+            # don't need that because they are closed on their destruction
  151
+            # automatically.
  152
+            try:
  153
+                for item in gen:
  154
+                    yield item
  155
+            finally:
  156
+                if hasattr(gen, 'close'):
  157
+                    gen.close()
  158
+
  159
+    # The trick is to start the generator.  Then the code execution runs until
  160
+    # the first dummy None is yielded at which point the context was already
  161
+    # pushed.  This item is discarded.  Then when the iteration continues the
  162
+    # real generator is executed.
  163
+    wrapped_g = generator()
  164
+    wrapped_g.next()
  165
+    return wrapped_g
  166
+
  167
+
95 168
 def jsonify(*args, **kwargs):
96 169
     """Creates a :class:`~flask.Response` with the JSON representation of
97 170
     the given arguments with an `application/json` mimetype.  The arguments
20  flask/testsuite/appctx.py
@@ -75,6 +75,26 @@ def __init__(self):
75 75
             self.assert_equal(
76 76
                 flask.render_template_string('{{ g.spam }}'), 'eggs')
77 77
 
  78
+    def test_context_refcounts(self):
  79
+        called = []
  80
+        app = flask.Flask(__name__)
  81
+        @app.teardown_request
  82
+        def teardown_req(error=None):
  83
+            called.append('request')
  84
+        @app.teardown_appcontext
  85
+        def teardown_app(error=None):
  86
+            called.append('app')
  87
+        @app.route('/')
  88
+        def index():
  89
+            with flask._app_ctx_stack.top:
  90
+                with flask._request_ctx_stack.top:
  91
+                    pass
  92
+            self.assert_(flask._request_ctx_stack.request.environ
  93
+                ['werkzeug.request'] is not None)
  94
+        c = app.test_client()
  95
+        c.get('/')
  96
+        self.assertEqual(called, ['request', 'app'])
  97
+
78 98
 
79 99
 def suite():
80 100
     suite = unittest.TestSuite()
59  flask/testsuite/helpers.py
@@ -397,6 +397,64 @@ def test_name_with_import_error(self):
397 397
             self.fail('Flask(import_name) is importing import_name.')
398 398
 
399 399
 
  400
+class StreamingTestCase(FlaskTestCase):
  401
+
  402
+    def test_streaming_with_context(self):
  403
+        app = flask.Flask(__name__)
  404
+        app.testing = True
  405
+        @app.route('/')
  406
+        def index():
  407
+            def generate():
  408
+                yield 'Hello '
  409
+                yield flask.request.args['name']
  410
+                yield '!'
  411
+            return flask.Response(flask.stream_with_context(generate()))
  412
+        c = app.test_client()
  413
+        rv = c.get('/?name=World')
  414
+        self.assertEqual(rv.data, 'Hello World!')
  415
+
  416
+    def test_streaming_with_context_as_decorator(self):
  417
+        app = flask.Flask(__name__)
  418
+        app.testing = True
  419
+        @app.route('/')
  420
+        def index():
  421
+            @flask.stream_with_context
  422
+            def generate():
  423
+                yield 'Hello '
  424
+                yield flask.request.args['name']
  425
+                yield '!'
  426
+            return flask.Response(generate())
  427
+        c = app.test_client()
  428
+        rv = c.get('/?name=World')
  429
+        self.assertEqual(rv.data, 'Hello World!')
  430
+
  431
+    def test_streaming_with_context_and_custom_close(self):
  432
+        app = flask.Flask(__name__)
  433
+        app.testing = True
  434
+        called = []
  435
+        class Wrapper(object):
  436
+            def __init__(self, gen):
  437
+                self._gen = gen
  438
+            def __iter__(self):
  439
+                return self
  440
+            def close(self):
  441
+                called.append(42)
  442
+            def next(self):
  443
+                return self._gen.next()
  444
+        @app.route('/')
  445
+        def index():
  446
+            def generate():
  447
+                yield 'Hello '
  448
+                yield flask.request.args['name']
  449
+                yield '!'
  450
+            return flask.Response(flask.stream_with_context(
  451
+                Wrapper(generate())))
  452
+        c = app.test_client()
  453
+        rv = c.get('/?name=World')
  454
+        self.assertEqual(rv.data, 'Hello World!')
  455
+        self.assertEqual(called, [42])
  456
+
  457
+
400 458
 def suite():
401 459
     suite = unittest.TestSuite()
402 460
     if flask.json_available:
@@ -404,4 +462,5 @@ def suite():
404 462
     suite.addTest(unittest.makeSuite(SendfileTestCase))
405 463
     suite.addTest(unittest.makeSuite(LoggingTestCase))
406 464
     suite.addTest(unittest.makeSuite(NoImportsTestCase))
  465
+    suite.addTest(unittest.makeSuite(StreamingTestCase))
407 466
     return suite

0 notes on commit d521899

Please sign in to comment.
Something went wrong with that request. Please try again.