new debug/test preserve context implementation

Author: davidism

CHANGES.rst

@@ -42,6 +42,18 @@ Unreleased
     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`
+-   A new implementation of context preservation during debugging and
+    testing. :pr:`4666`
+
+    -   ``request``, ``g``, and other context-locals point to the
+        correct data when running code in the interactive debugger
+        console. :issue:`2836`
+    -   Teardown functions are always run at the end of the request,
+        even if the context is preserved. They are also run after the
+        preserved context is popped.
+    -   ``stream_with_context`` preserves context separately from a
+        ``with client`` block. It will be cleaned up when
+        ``response.get_data()`` or ``response.close()`` is called.
 
 
 Version 2.1.3

docs/config.rst

@@ -126,14 +126,6 @@ The following configuration values are used internally by Flask:
 
     Default: ``None``
 
-.. py:data:: PRESERVE_CONTEXT_ON_EXCEPTION
-
-    Don't pop the request context when an exception occurs. If not set, this
-    is true if ``DEBUG`` is true. This allows debuggers to introspect the
-    request data on errors, and should normally not need to be set directly.
-
-    Default: ``None``
-
 .. py:data:: TRAP_HTTP_EXCEPTIONS
 
     If there is no handler for an ``HTTPException``-type exception, re-raise it
@@ -392,6 +384,9 @@ The following configuration values are used internally by Flask:
 
     Added :data:`MAX_COOKIE_SIZE` to control a warning from Werkzeug.
 
+.. versionchanged:: 2.2
+    Removed ``PRESERVE_CONTEXT_ON_EXCEPTION``.
+
 
 Configuring from Python Files
 -----------------------------

docs/reqcontext.rst

@@ -219,25 +219,6 @@ sent:
     :meth:`~Flask.teardown_request` functions are called.
 
 
-Context Preservation on Error
------------------------------
-
-At the end of a request, the request context is popped and all data
-associated with it is destroyed. If an error occurs during development,
-it is useful to delay destroying the data for debugging purposes.
-
-When the development server is running in development mode (the
-``--env`` option is set to ``'development'``), the error and data will
-be preserved and shown in the interactive debugger.
-
-This behavior can be controlled with the
-:data:`PRESERVE_CONTEXT_ON_EXCEPTION` config. As described above, it
-defaults to ``True`` in the development environment.
-
-Do not enable :data:`PRESERVE_CONTEXT_ON_EXCEPTION` in production, as it
-will cause your application to leak memory on exceptions.
-
-
 .. _notes-on-proxies:
 
 Notes On Proxies

src/flask/app.py

@@ -331,7 +331,6 @@ class Flask(Scaffold):
             "DEBUG": None,
             "TESTING": False,
             "PROPAGATE_EXCEPTIONS": None,
-            "PRESERVE_CONTEXT_ON_EXCEPTION": None,
             "SECRET_KEY": None,
             "PERMANENT_SESSION_LIFETIME": timedelta(days=31),
             "USE_X_SENDFILE": False,
@@ -583,19 +582,6 @@ def propagate_exceptions(self) -> bool:
             return rv
         return self.testing or self.debug
 
-    @property
-    def preserve_context_on_exception(self) -> bool:
-        """Returns the value of the ``PRESERVE_CONTEXT_ON_EXCEPTION``
-        configuration value in case it's set, otherwise a sensible default
-        is returned.
-
-        .. versionadded:: 0.7
-        """
-        rv = self.config["PRESERVE_CONTEXT_ON_EXCEPTION"]
-        if rv is not None:
-            return rv
-        return self.debug
-
     @locked_cached_property
     def logger(self) -> logging.Logger:
         """A standard Python :class:`~logging.Logger` for the app, with
@@ -2301,9 +2287,14 @@ def wsgi_app(self, environ: dict, start_response: t.Callable) -> t.Any:
                 raise
             return response(environ, start_response)
         finally:
-            if self.should_ignore_error(error):
+            if "werkzeug.debug.preserve_context" in environ:
+                environ["werkzeug.debug.preserve_context"](_app_ctx_stack.top)
+                environ["werkzeug.debug.preserve_context"](_request_ctx_stack.top)
+
+            if error is not None and self.should_ignore_error(error):
                 error = None
-            ctx.auto_pop(error)
+
+            ctx.pop(error)
 
     def __call__(self, environ: dict, start_response: t.Callable) -> t.Any:
         """The WSGI server calls the Flask application object as the

src/flask/ctx.py

@@ -289,20 +289,12 @@ class RequestContext:
     functions registered on the application for teardown execution
     (:meth:`~flask.Flask.teardown_request`).
 
-    The request context is automatically popped at the end of the request
-    for you.  In debug mode the request context is kept around if
-    exceptions happen so that interactive debuggers have a chance to
-    introspect the data.  With 0.4 this can also be forced for requests
-    that did not fail and outside of ``DEBUG`` mode.  By setting
-    ``'flask._preserve_context'`` to ``True`` on the WSGI environment the
-    context will not pop itself at the end of the request.  This is used by
-    the :meth:`~flask.Flask.test_client` for example to implement the
-    deferred cleanup functionality.
-
-    You might find this helpful for unittests where you need the
-    information from the context local around for a little longer.  Make
-    sure to properly :meth:`~werkzeug.LocalStack.pop` the stack yourself in
-    that situation, otherwise your unittests will leak memory.
+    The request context is automatically popped at the end of the
+    request. When using the interactive debugger, the context will be
+    restored so ``request`` is still accessible. Similarly, the test
+    client can preserve the context after the request ends. However,
+    teardown functions may already have closed some resources such as
+    database connections.
     """
 
     def __init__(
@@ -330,14 +322,6 @@ def __init__(
         # one is created implicitly so for each level we add this information
         self._implicit_app_ctx_stack: t.List[t.Optional["AppContext"]] = []
 
-        # indicator if the context was preserved.  Next time another context
-        # is pushed the preserved context is popped.
-        self.preserved = False
-
-        # remembers the exception for pop if there is one in case the context
-        # preservation kicks in.
-        self._preserved_exc = None
-
         # Functions that should be executed after the request on the response
         # object.  These will be called before the regular "after_request"
         # functions.
@@ -400,19 +384,6 @@ def match_request(self) -> None:
             self.request.routing_exception = e
 
     def push(self) -> None:
-        """Binds the request context to the current context."""
-        # If an exception occurs in debug mode or if context preservation is
-        # activated under exception situations exactly one context stays
-        # on the stack.  The rationale is that you want to access that
-        # information under debug situations.  However if someone forgets to
-        # pop that context again we want to make sure that on the next push
-        # it's invalidated, otherwise we run at risk that something leaks
-        # memory.  This is usually only a problem in test suite since this
-        # functionality is not active in production environments.
-        top = _request_ctx_stack.top
-        if top is not None and top.preserved:
-            top.pop(top._preserved_exc)
-
         # Before we push the request context we have to ensure that there
         # is an application context.
         app_ctx = _app_ctx_stack.top
@@ -454,8 +425,6 @@ def pop(self, exc: t.Optional[BaseException] = _sentinel) -> None:  # type: igno
 
         try:
             if not self._implicit_app_ctx_stack:
-                self.preserved = False
-                self._preserved_exc = None
                 if exc is _sentinel:
                     exc = sys.exc_info()[1]
                 self.app.do_teardown_request(exc)
@@ -481,13 +450,18 @@ def pop(self, exc: t.Optional[BaseException] = _sentinel) -> None:  # type: igno
             ), f"Popped wrong request context. ({rv!r} instead of {self!r})"
 
     def auto_pop(self, exc: t.Optional[BaseException]) -> None:
-        if self.request.environ.get("flask._preserve_context") or (
-            exc is not None and self.app.preserve_context_on_exception
-        ):
-            self.preserved = True
-            self._preserved_exc = exc  # type: ignore
-        else:
-            self.pop(exc)
+        """
+        .. deprecated:: 2.2
+            Will be removed in Flask 2.3.
+        """
+        import warnings
+
+        warnings.warn(
+            "'ctx.auto_pop' is deprecated and will be removed in Flask 2.3.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self.pop(exc)
 
     def __enter__(self) -> "RequestContext":
         self.push()
@@ -499,12 +473,7 @@ def __exit__(
         exc_value: t.Optional[BaseException],
         tb: t.Optional[TracebackType],
     ) -> None:
-        # do not pop the request stack if we are in debug mode and an
-        # exception happened.  This will allow the debugger to still
-        # access the request object in the interactive shell.  Furthermore
-        # the context can be force kept alive for the test client.
-        # See flask.testing for how this works.
-        self.auto_pop(exc_value)
+        self.pop(exc_value)
 
     def __repr__(self) -> str:
         return (

src/flask/scaffold.py

@@ -600,13 +600,6 @@ def teardown_request(self, f: ft.TeardownCallable) -> ft.TeardownCallable:
         be passed an error object.
 
         The return values of teardown functions are ignored.
-
-        .. admonition:: Debug Note
-
-           In debug mode Flask will not tear down a request on an exception
-           immediately.  Instead it will keep it alive so that the interactive
-           debugger can still access it.  This behavior can be controlled
-           by the ``PRESERVE_CONTEXT_ON_EXCEPTION`` configuration variable.
         """
         self.teardown_request_funcs.setdefault(None, []).append(f)
         return f

src/flask/testing.py

@@ -1,5 +1,6 @@
 import typing as t
 from contextlib import contextmanager
+from contextlib import ExitStack
 from copy import copy
 from types import TracebackType
 
@@ -108,10 +109,12 @@ class FlaskClient(Client):
     """
 
     application: "Flask"
-    preserve_context = False
 
     def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
         super().__init__(*args, **kwargs)
+        self.preserve_context = False
+        self._new_contexts: t.List[t.ContextManager[t.Any]] = []
+        self._context_stack = ExitStack()
         self.environ_base = {
             "REMOTE_ADDR": "127.0.0.1",
             "HTTP_USER_AGENT": f"werkzeug/{werkzeug.__version__}",
@@ -173,11 +176,12 @@ def session_transaction(
             self.cookie_jar.extract_wsgi(c.request.environ, headers)
 
     def _copy_environ(self, other):
-        return {
-            **self.environ_base,
-            **other,
-            "flask._preserve_context": self.preserve_context,
-        }
+        out = {**self.environ_base, **other}
+
+        if self.preserve_context:
+            out["werkzeug.debug.preserve_context"] = self._new_contexts.append
+
+        return out
 
     def _request_from_builder_args(self, args, kwargs):
         kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {}))
@@ -214,12 +218,24 @@ def open(
             # request is None
             request = self._request_from_builder_args(args, kwargs)
 
-        return super().open(
+        # Pop any previously preserved contexts. This prevents contexts
+        # from being preserved across redirects or multiple requests
+        # within a single block.
+        self._context_stack.close()
+
+        response = super().open(
             request,
             buffered=buffered,
             follow_redirects=follow_redirects,
         )
 
+        # Re-push contexts that were preserved during the request.
+        while self._new_contexts:
+            cm = self._new_contexts.pop()
+            self._context_stack.enter_context(cm)
+
+        return response
+
     def __enter__(self) -> "FlaskClient":
         if self.preserve_context:
             raise RuntimeError("Cannot nest client invocations")
@@ -233,18 +249,7 @@ def __exit__(
         tb: t.Optional[TracebackType],
     ) -> None:
         self.preserve_context = False
-
-        # Normally the request context is preserved until the next
-        # request in the same thread comes. When the client exits we
-        # want to clean up earlier. Pop request contexts until the stack
-        # is empty or a non-preserved one is found.
-        while True:
-            top = _request_ctx_stack.top
-
-            if top is not None and top.preserved:
-                top.pop()
-            else:
-                break
+        self._context_stack.close()
 
 
 class FlaskCliRunner(CliRunner):

tests/test_basic.py

@@ -928,13 +928,8 @@ def test_baseexception_error_handling(app, client):
     def broken_func():
         raise KeyboardInterrupt()
 
-    with client:
-        with pytest.raises(KeyboardInterrupt):
-            client.get("/")
-
-        ctx = flask._request_ctx_stack.top
-        assert ctx.preserved
-        assert type(ctx._preserved_exc) is KeyboardInterrupt
+    with pytest.raises(KeyboardInterrupt):
+        client.get("/")
 
 
 def test_before_request_and_routing_errors(app, client):
@@ -1769,57 +1764,6 @@ def for_bar_foo():
     assert client.get("/bar/123").data == b"123"
 
 
-def test_preserve_only_once(app, client):
-    app.debug = True
-
-    @app.route("/fail")
-    def fail_func():
-        1 // 0
-
-    for _x in range(3):
-        with pytest.raises(ZeroDivisionError):
-            client.get("/fail")
-
-    assert flask._request_ctx_stack.top is not None
-    assert flask._app_ctx_stack.top is not None
-    # implicit appctx disappears too
-    flask._request_ctx_stack.top.pop()
-    assert flask._request_ctx_stack.top is None
-    assert flask._app_ctx_stack.top is None
-
-
-def test_preserve_remembers_exception(app, client):
-    app.debug = True
-    errors = []
-
-    @app.route("/fail")
-    def fail_func():
-        1 // 0
-
-    @app.route("/success")
-    def success_func():
-        return "Okay"
-
-    @app.teardown_request
-    def teardown_handler(exc):
-        errors.append(exc)
-
-    # After this failure we did not yet call the teardown handler
-    with pytest.raises(ZeroDivisionError):
-        client.get("/fail")
-    assert errors == []
-
-    # But this request triggers it, and it's an error
-    client.get("/success")
-    assert len(errors) == 2
-    assert isinstance(errors[0], ZeroDivisionError)
-
-    # At this point another request does nothing.
-    client.get("/success")
-    assert len(errors) == 3
-    assert errors[1] is None
-
-
 def test_get_method_on_g(app_ctx):
     assert flask.g.get("x") is None
     assert flask.g.get("x", 11) == 11