Merge pull request #4629 from pgjones/generate

allow view to return generator

CHANGES.rst

@@ -38,6 +38,10 @@ Unreleased
     context will already be active at that point. :issue:`2410`
 -   ``SessionInterface.get_expiration_time`` uses a timezone-aware
     value. :pr:`4645`
+-   View functions can return generators directly instead of wrapping
+    them in a ``Response``. :pr:`4629`
+-   Add ``stream_template`` and ``stream_template_string`` functions to
+    render a template as a stream of pieces. :pr:`4629`
 
 
 Version 2.1.3

docs/api.rst

@@ -287,6 +287,10 @@ Template Rendering
 
 .. autofunction:: render_template_string
 
+.. autofunction:: stream_template
+
+.. autofunction:: stream_template_string
+
 .. autofunction:: get_template_attribute
 
 Configuration

docs/patterns/streaming.rst

@@ -20,7 +20,7 @@ data and to then invoke that function and pass it to a response object::
         def generate():
             for row in iter_all_rows():
                 yield f"{','.join(row)}\n"
-        return app.response_class(generate(), mimetype='text/csv')
+        return generate(), {"Content-Type": "text/csv")
 
 Each ``yield`` expression is directly sent to the browser.  Note though
 that some WSGI middlewares might break streaming, so be careful there in
@@ -29,52 +29,57 @@ debug environments with profilers and other things you might have enabled.
 Streaming from Templates
 ------------------------
 
-The Jinja2 template engine also supports rendering templates piece by
+The Jinja2 template engine supports rendering a template piece by
-piece.  This functionality is not directly exposed by Flask because it is
+piece, returning an iterator of strings. Flask provides the
-quite uncommon, but you can easily do it yourself::
+:func:`~flask.stream_template` and :func:`~flask.stream_template_string`
+functions to make this easier to use.
 
-    def stream_template(template_name, **context):
+.. code-block:: python
-        app.update_template_context(context)
-        t = app.jinja_env.get_template(template_name)
-        rv = t.stream(context)
-        rv.enable_buffering(5)
-        return rv
 
-    @app.route('/my-large-page.html')
+    from flask import stream_template
-    def render_large_template():
+
-        rows = iter_all_rows()
+    @app.get("/timeline")
-        return app.response_class(stream_template('the_template.html', rows=rows))
+    def timeline():
+        return stream_template("timeline.html")
+
+The parts yielded by the render stream tend to match statement blocks in
+the template.
 
-The trick here is to get the template object from the Jinja2 environment
-on the application and to call :meth:`~jinja2.Template.stream` instead of
-:meth:`~jinja2.Template.render` which returns a stream object instead of a
-string.  Since we're bypassing the Flask template render functions and
-using the template object itself we have to make sure to update the render
-context ourselves by calling :meth:`~flask.Flask.update_template_context`.
-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
+The :data:`~flask.request` will not be active while the generator is
+running, because the view has already returned at that point. If you try
+to access ``request``, you'll get a ``RuntimeError``.
 
-Note that when you stream data, the request context is already gone the
+If your generator function relies on data in ``request``, use the
-moment the function executes.  Flask 0.9 provides you with a helper that
+:func:`~flask.stream_with_context` wrapper. This will keep the request
-can keep the request context around during the execution of the
+context active during the generator.
-generator::
+
+.. code-block:: python
 
     from flask import stream_with_context, request
+    from markupsafe import escape
 
     @app.route('/stream')
     def streamed_response():
         def generate():
-            yield 'Hello '
+            yield '<p>Hello '
-            yield request.args['name']
+            yield escape(request.args['name'])
-            yield '!'
+            yield '!</p>'
-        return app.response_class(stream_with_context(generate()))
+        return stream_with_context(generate())
 
-Without the :func:`~flask.stream_with_context` function you would get a
+It can also be used as a decorator.
-:class:`RuntimeError` at that point.
+
+.. code-block:: python
+
+    @stream_with_context
+    def generate():
+        ...
+
+    return generate()
+
+The :func:`~flask.stream_template` and
+:func:`~flask.stream_template_string` functions automatically
+use :func:`~flask.stream_with_context` if a request is active.

docs/templating.rst

@@ -201,3 +201,29 @@ templates::
 You could also build `format_price` as a template filter (see
 :ref:`registering-filters`), but this demonstrates how to pass functions in a
 context processor.
+
+Streaming
+---------
+
+It can be useful to not render the whole template as one complete
+string, instead render it as a stream, yielding smaller incremental
+strings. This can be used for streaming HTML in chunks to speed up
+initial page load, or to save memory when rendering a very large
+template.
+
+The Jinja2 template engine supports rendering a template piece
+by piece, returning an iterator of strings. Flask provides the
+:func:`~flask.stream_template` and :func:`~flask.stream_template_string`
+functions to make this easier to use.
+
+.. code-block:: python
+
+    from flask import stream_template
+
+    @app.get("/timeline")
+    def timeline():
+        return stream_template("timeline.html")
+
+These functions automatically apply the
+:func:`~flask.stream_with_context` wrapper if a request is active, so
+that it remains available in the template.

src/flask/__init__.py

@@ -41,5 +41,7 @@ from .signals import signals_available as signals_available
 from .signals import template_rendered as template_rendered
 from .templating import render_template as render_template
 from .templating import render_template_string as render_template_string
+from .templating import stream_template as stream_template
+from .templating import stream_template_string as stream_template_string
 
 __version__ = "2.2.0.dev0"

src/flask/app.py

@@ -5,6 +5,7 @@ import os
 import sys
 import typing as t
 import weakref
+from collections.abc import Iterator as _abc_Iterator
 from datetime import timedelta
 from itertools import chain
 from threading import Lock
@@ -1843,6 +1844,10 @@ class Flask(Scaffold):
             ``dict``
                 A dictionary that will be jsonify'd before being returned.
 
+            ``generator`` or ``iterator``
+                A generator that returns ``str`` or ``bytes`` to be
+                streamed as the response.
+
             ``tuple``
                 Either ``(body, status, headers)``, ``(body, status)``, or
                 ``(body, headers)``, where ``body`` is any of the other types
@@ -1862,6 +1867,12 @@ class Flask(Scaffold):
                 The function is called as a WSGI application. The result is
                 used to create a response object.
 
+        .. versionchanged:: 2.2
+            A generator will be converted to a streaming response.
+
+        .. versionchanged:: 1.1
+            A dict will be converted to a JSON response.
+
         .. versionchanged:: 0.9
            Previously a tuple was interpreted as the arguments for the
            response object.
@@ -1900,7 +1911,7 @@ class Flask(Scaffold):
 
         # make sure the body is an instance of the response class
         if not isinstance(rv, self.response_class):
-            if isinstance(rv, (str, bytes, bytearray)):
+            if isinstance(rv, (str, bytes, bytearray)) or isinstance(rv, _abc_Iterator):
                 # let the response class set the status and headers instead of
                 # waiting to do it manually, so that the class can handle any
                 # special logic

src/flask/templating.py

@@ -7,6 +7,9 @@ from jinja2 import TemplateNotFound
 
 from .globals import _app_ctx_stack
 from .globals import _request_ctx_stack
+from .globals import current_app
+from .globals import request
+from .helpers import stream_with_context
 from .signals import before_render_template
 from .signals import template_rendered
 
@@ -122,8 +125,6 @@ class DispatchingJinjaLoader(BaseLoader):
 
 
 def _render(template: Template, context: dict, app: "Flask") -> str:
-    """Renders the template and fires the signal"""
-
     before_render_template.send(app, template=template, context=context)
     rv = template.render(context)
     template_rendered.send(app, template=template, context=context)
@@ -164,3 +165,56 @@ def render_template_string(source: str, **context: t.Any) -> str:
     ctx = _app_ctx_stack.top
     ctx.app.update_template_context(context)
     return _render(ctx.app.jinja_env.from_string(source), context, ctx.app)
+
+
+def _stream(
+    app: "Flask", template: Template, context: t.Dict[str, t.Any]
+) -> t.Iterator[str]:
+    app.update_template_context(context)
+    before_render_template.send(app, template=template, context=context)
+
+    def generate() -> t.Iterator[str]:
+        yield from template.generate(context)
+        template_rendered.send(app, template=template, context=context)
+
+    rv = generate()
+
+    # If a request context is active, keep it while generating.
+    if request:
+        rv = stream_with_context(rv)
+
+    return rv
+
+
+def stream_template(
+    template_name_or_list: t.Union[str, Template, t.List[t.Union[str, Template]]],
+    **context: t.Any
+) -> t.Iterator[str]:
+    """Render a template by name with the given context as a stream.
+    This returns an iterator of strings, which can be used as a
+    streaming response from a view.
+
+    :param template_name_or_list: The name of the template to render. If
+        a list is given, the first name to exist will be rendered.
+    :param context: The variables to make available in the template.
+
+    .. versionadded:: 2.2
+    """
+    app = current_app._get_current_object()  # type: ignore[attr-defined]
+    template = app.jinja_env.get_or_select_template(template_name_or_list)
+    return _stream(app, template, context)
+
+
+def stream_template_string(source: str, **context: t.Any) -> t.Iterator[str]:
+    """Render a template from the given source string with the given
+    context as a stream. This returns an iterator of strings, which can
+    be used as a streaming response from a view.
+
+    :param source: The source code of the template to render.
+    :param context: The variables to make available in the template.
+
+    .. versionadded:: 2.2
+    """
+    app = current_app._get_current_object()  # type: ignore[attr-defined]
+    template = app.jinja_env.from_string(source)
+    return _stream(app, template, context)

src/flask/typing.py

@@ -6,7 +6,9 @@ if t.TYPE_CHECKING:  # pragma: no cover
     from werkzeug.wrappers import Response  # noqa: F401
 
 # The possible types that are directly convertible or are a Response object.
-ResponseValue = t.Union["Response", str, bytes, t.Dict[str, t.Any]]
+ResponseValue = t.Union[
+    "Response", str, bytes, t.Dict[str, t.Any], t.Iterator[str], t.Iterator[bytes]
+]
 
 # the possible types for an individual HTTP header
 # This should be a Union, but mypy doesn't pass unless it's a TypeVar.

tests/test_basic.py

@@ -1276,6 +1276,11 @@ def test_make_response(app, req_ctx):
     assert rv.data == b"W00t"
     assert rv.mimetype == "text/html"
 
+    rv = flask.make_response(c for c in "Hello")
+    assert rv.status_code == 200
+    assert rv.data == b"Hello"
+    assert rv.mimetype == "text/html"
+
 
 def test_make_response_with_response_instance(app, req_ctx):
     rv = flask.make_response(flask.jsonify({"msg": "W00t"}), 400)

tests/test_templating.py

@@ -29,6 +29,15 @@ def test_original_win(app, client):
     assert rv.data == b"42"
 
 
+def test_simple_stream(app, client):
+    @app.route("/")
+    def index():
+        return flask.stream_template_string("{{ config }}", config=42)
+
+    rv = client.get("/")
+    assert rv.data == b"42"
+
+
 def test_request_less_rendering(app, app_ctx):
     app.config["WORLD_NAME"] = "Special World"
 

tests/typing/typing_route.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import typing as t
 from http import HTTPStatus
 
 from flask import Flask
 from flask import jsonify
+from flask import stream_template
 from flask.templating import render_template
 from flask.views import View
 from flask.wrappers import Response
@@ -26,6 +28,25 @@ def hello_json() -> Response:
     return jsonify({"response": "Hello, World!"})
 
 
+@app.route("/generator")
+def hello_generator() -> t.Generator[str, None, None]:
+    def show() -> t.Generator[str, None, None]:
+        for x in range(100):
+            yield f"data:{x}\n\n"
+
+    return show()
+
+
+@app.route("/generator-expression")
+def hello_generator_expression() -> t.Iterator[bytes]:
+    return (f"data:{x}\n\n".encode() for x in range(100))
+
+
+@app.route("/iterator")
+def hello_iterator() -> t.Iterator[str]:
+    return iter([f"data:{x}\n\n" for x in range(100)])
+
+
 @app.route("/status")
 @app.route("/status/<int:code>")
 def tuple_status(code: int = 200) -> tuple[str, int]:
@@ -48,6 +69,11 @@ def return_template(name: str | None = None) -> str:
     return render_template("index.html", name=name)
 
 
+@app.route("/template")
+def return_template_stream() -> t.Iterator[str]:
+    return stream_template("index.html", name="Hello")
+
+
 class RenderTemplateView(View):
     def __init__(self: RenderTemplateView, template_name: str) -> None:
         self.template_name = template_name