Permalink
Browse files

Added flask.stream_with_context

  • Loading branch information...
1 parent 2e816f5 commit d5218997d927be869dd55ef04542e1bbc1e69653 @mitsuhiko mitsuhiko committed Jun 27, 2012
Showing with 222 additions and 27 deletions.
  1. +2 −0 CHANGES
  2. +5 −0 docs/api.rst
  3. +23 −0 docs/patterns/streaming.rst
  4. +2 −1 flask/__init__.py
  5. +38 −26 flask/ctx.py
  6. +73 −0 flask/helpers.py
  7. +20 −0 flask/testsuite/appctx.py
  8. +59 −0 flask/testsuite/helpers.py
View
@@ -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
-------------
View
@@ -375,6 +375,11 @@ Extensions
.. versionadded:: 0.8
+Stream Helpers
+--------------
+
+.. autofunction:: stream_with_context
+
Useful Internals
----------------
@@ -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.
View
@@ -22,7 +22,8 @@
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, \
View
@@ -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,15 +102,22 @@ def __init__(self, app):
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."""
- if exc is None:
- exc = sys.exc_info()[1]
- self.app.do_teardown_appcontext(exc)
+ self._refcnt -= 1
+ if self._refcnt <= 0:
+ if exc is None:
+ exc = sys.exc_info()[1]
+ self.app.do_teardown_appcontext(exc)
rv = _app_ctx_stack.pop()
assert rv is self, 'Popped wrong app context. (%r instead of %r)' \
% (rv, self)
@@ -128,7 +127,7 @@ def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, tb):
- self.pop()
+ self.pop(exc_value)
class RequestContext(object):
@@ -169,15 +168,16 @@ def __init__(self, app, environ):
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 @@ def push(self):
# 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 @@ def pop(self, exc=None):
.. versionchanged:: 0.9
Added the `exc` argument.
"""
- self.preserved = False
- if exc is None:
- exc = sys.exc_info()[1]
- self.app.do_teardown_request(exc)
+ 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.
- rv.request.environ['werkzeug.request'] = None
+ 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()
View
@@ -21,6 +21,7 @@
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
@@ -75,6 +75,26 @@ def __init__(self):
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()
@@ -397,11 +397,70 @@ def test_name_with_import_error(self):
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:
suite.addTest(unittest.makeSuite(JSONTestCase))
suite.addTest(unittest.makeSuite(SendfileTestCase))
suite.addTest(unittest.makeSuite(LoggingTestCase))
suite.addTest(unittest.makeSuite(NoImportsTestCase))
+ suite.addTest(unittest.makeSuite(StreamingTestCase))
return suite

0 comments on commit d521899

Please sign in to comment.