diff --git a/CHANGES b/CHANGES index 28a0790c..acdfee69 100644 --- a/CHANGES +++ b/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 ------------- diff --git a/docs/api.rst b/docs/api.rst index dcb54baf..8a7b5ce0 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -375,6 +375,11 @@ Extensions .. versionadded:: 0.8 +Stream Helpers +-------------- + +.. autofunction:: stream_with_context + Useful Internals ---------------- diff --git a/docs/patterns/streaming.rst b/docs/patterns/streaming.rst index 8393b00b..ac232dcc 100644 --- a/docs/patterns/streaming.rst +++ b/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. diff --git a/flask/__init__.py b/flask/__init__.py index de84bb69..e48f7a97 100644 --- a/flask/__init__.py +++ b/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, \ diff --git a/flask/ctx.py b/flask/ctx.py index 0cf34491..3ea42a27 100644 --- a/flask/ctx.py +++ b/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,15 +102,22 @@ 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.""" - 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 @@ 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. """ - 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() diff --git a/flask/helpers.py b/flask/helpers.py index 631e29be..501a2f81 100644 --- a/flask/helpers.py +++ b/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 diff --git a/flask/testsuite/appctx.py b/flask/testsuite/appctx.py index 1dcdb406..6454389e 100644 --- a/flask/testsuite/appctx.py +++ b/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() diff --git a/flask/testsuite/helpers.py b/flask/testsuite/helpers.py index 816f6cd8..54c01482 100644 --- a/flask/testsuite/helpers.py +++ b/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