From 9822a0351574790cb66c652fcc396ad7aa2b09d8 Mon Sep 17 00:00:00 2001 From: David Lord Date: Tue, 19 Aug 2025 08:18:55 -0700 Subject: [PATCH] refactor stream_with_context for async views --- CHANGES.rst | 1 + src/flask/helpers.py | 69 ++++++++++++++++++++++++------------------- tests/test_helpers.py | 23 +++++++++++++++ 3 files changed, 62 insertions(+), 31 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index a5761bc2..b8625dc1 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -3,6 +3,7 @@ Version 3.1.2 Unreleased +- ``stream_with_context`` does not fail inside async views. :issue:`5774` - When using ``follow_redirects`` in the test client, the final state of ``session`` is correct. :issue:`5786` - Relax type hint for passing bytes IO to ``send_file``. :issue:`5776` diff --git a/src/flask/helpers.py b/src/flask/helpers.py index 081548a9..5d412c90 100644 --- a/src/flask/helpers.py +++ b/src/flask/helpers.py @@ -13,6 +13,7 @@ from werkzeug.exceptions import abort as _wz_abort from werkzeug.utils import redirect as _wz_redirect from werkzeug.wrappers import Response as BaseResponse +from .globals import _cv_app from .globals import _cv_request from .globals import current_app from .globals import request @@ -62,35 +63,40 @@ def stream_with_context( def stream_with_context( generator_or_function: t.Iterator[t.AnyStr] | t.Callable[..., t.Iterator[t.AnyStr]], ) -> t.Iterator[t.AnyStr] | t.Callable[[t.Iterator[t.AnyStr]], t.Iterator[t.AnyStr]]: - """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. + """Wrap a response generator function so that it runs inside the current + request context. This keeps :data:`request`, :data:`session`, and :data:`g` + available, even though at the point the generator runs the request context + will typically have ended. - This function however can help you keep the context around for longer:: + Use it as a decorator on a generator function: + + .. code-block:: python from flask import stream_with_context, request, Response - @app.route('/stream') + @app.get("/stream") def streamed_response(): @stream_with_context def generate(): - yield 'Hello ' - yield request.args['name'] - yield '!' + yield "Hello " + yield request.args["name"] + yield "!" + return Response(generate()) - Alternatively it can also be used around a specific generator:: + Or use it as a wrapper around a created generator: + + .. code-block:: python from flask import stream_with_context, request, Response - @app.route('/stream') + @app.get("/stream") def streamed_response(): def generate(): - yield 'Hello ' - yield request.args['name'] - yield '!' + yield "Hello " + yield request.args["name"] + yield "!" + return Response(stream_with_context(generate())) .. versionadded:: 0.9 @@ -105,35 +111,36 @@ def stream_with_context( return update_wrapper(decorator, generator_or_function) # type: ignore[arg-type] - def generator() -> t.Iterator[t.AnyStr | None]: - ctx = _cv_request.get(None) - if ctx is None: + def generator() -> t.Iterator[t.AnyStr]: + if (req_ctx := _cv_request.get(None)) is None: raise RuntimeError( "'stream_with_context' can only be used when a request" " context is active, such as in a view function." ) - 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. + app_ctx = _cv_app.get() + # Setup code below will run the generator to this point, so that the + # current contexts are recorded. The contexts must be pushed after, + # otherwise their ContextVar will record the wrong event loop during + # async view functions. + yield None # type: ignore[misc] + + # Push the app context first, so that the request context does not + # automatically create and push a different app context. + with app_ctx, req_ctx: try: yield from gen finally: + # Clean up in case the user wrapped a WSGI iterator. 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. + # Execute the generator to the sentinel value. This ensures the context is + # preserved in the generator's state. Further iteration will push the + # context and yield from the original iterator. wrapped_g = generator() next(wrapped_g) - return wrapped_g # type: ignore[return-value] + return wrapped_g def make_response(*args: t.Any) -> Response: diff --git a/tests/test_helpers.py b/tests/test_helpers.py index ee77f176..efd22aa9 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -306,6 +306,29 @@ class TestStreaming: rv = client.get("/") assert rv.data == b"flask" + def test_async_view(self, app, client): + @app.route("/") + async def index(): + flask.session["test"] = "flask" + + @flask.stream_with_context + def gen(): + yield flask.session["test"] + + return flask.Response(gen()) + + # response is closed without reading stream + client.get().close() + # response stream is read + assert client.get().text == "flask" + + # same as above, but with client context preservation + with client: + client.get().close() + + with client: + assert client.get().text == "flask" + class TestHelpers: @pytest.mark.parametrize(