Close #19403: make contextlib.redirect_stdout reentrant

4 files changed, 100 insertions, 55 deletions

diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
index 4bb54c5..82efd0c 100644
--- a/Doc/library/contextlib.rst
+++ b/Doc/library/contextlib.rst
@@ -651,22 +651,33 @@ managers can not only be used in multiple :keyword:`with` statements,
 but may also be used *inside* a :keyword:`with` statement that is already
 using the same context manager.
 
-:class:`threading.RLock` is an example of a reentrant context manager, as is
-:func:`suppress`. Here's a toy example of reentrant use (real world
-examples of reentrancy are more likely to occur with objects like recursive
-locks and are likely to be far more complicated than this example)::
-
-    >>> from contextlib import suppress
-    >>> ignore_raised_exception = suppress(ZeroDivisionError)
-    >>> with ignore_raised_exception:
-    ...     with ignore_raised_exception:
-    ...         1/0
-    ...     print("This line runs")
-    ...     1/0
-    ...     print("This is skipped")
+:class:`threading.RLock` is an example of a reentrant context manager, as are
+:func:`suppress` and :func:`redirect_stdout`. Here's a very simple example of
+reentrant use::
+
+    >>> from contextlib import redirect_stdout
+    >>> from io import StringIO
+    >>> stream = StringIO()
+    >>> write_to_stream = redirect_stdout(stream)
+    >>> with write_to_stream:
+    ...     print("This is written to the stream rather than stdout")
+    ...     with write_to_stream:
+    ...         print("This is also written to the stream")
     ...
-    This line runs
-    >>> # The second exception is also suppressed
+    >>> print("This is written directly to stdout")
+    This is written directly to stdout
+    >>> print(stream.getvalue())
+    This is written to the stream rather than stdout
+    This is also written to the stream
+
+Real world examples of reentrancy are more likely to involve multiple
+functions calling each other and hence be far more complicated than this
+example.
+
+Note also that being reentrant is *not* the same thing as being thread safe.
+:func:`redirect_stdout`, for example, is definitely not thread safe, as it
+makes a global modification to the system state by binding :data:`sys.stdout`
+to a different stream.
 
 
 .. _reusable-cms:
@@ -681,32 +692,58 @@ reusable). These context managers support being used multiple times, but
 will fail (or otherwise not work correctly) if the specific context manager
 instance has already been used in a containing with statement.
 
-An example of a reusable context manager is :func:`redirect_stdout`::
+:class:`threading.Lock` is an example of a reusable, but not reentrant,
+context manager (for a reentrant lock, it is necessary to use
+:class:`threading.RLock` instead).
 
-    >>> from contextlib import redirect_stdout
-    >>> from io import StringIO
-    >>> f = StringIO()
-    >>> collect_output = redirect_stdout(f)
-    >>> with collect_output:
-    ...     print("Collected")
+Another example of a reusable, but not reentrant, context manager is
+:class:`ExitStack`, as it invokes *all* currently registered callbacks
+when leaving any with statement, regardless of where those callbacks
+were added::
+
+    >>> from contextlib import ExitStack
+    >>> stack = ExitStack()
+    >>> with stack:
+    ...     stack.callback(print, "Callback: from first context")
+    ...     print("Leaving first context")
     ...
-    >>> print("Not collected")
-    Not collected
-    >>> with collect_output:
-    ...     print("Also collected")
+    Leaving first context
+    Callback: from first context
+    >>> with stack:
+    ...     stack.callback(print, "Callback: from second context")
+    ...     print("Leaving second context")
     ...
-    >>> print(f.getvalue())
-    Collected
-    Also collected
-
-However, this context manager is not reentrant, so attempting to reuse it
-within a containing with statement fails:
-
-    >>> with collect_output:
-    ...     # Nested reuse is not permitted
-    ...     with collect_output:
-    ...         pass
+    Leaving second context
+    Callback: from second context
+    >>> with stack:
+    ...     stack.callback(print, "Callback: from outer context")
+    ...     with stack:
+    ...         stack.callback(print, "Callback: from inner context")
+    ...         print("Leaving inner context")
+    ...     print("Leaving outer context")
     ...
-    Traceback (most recent call last):
-      ...
-    RuntimeError: Cannot reenter <...>
+    Leaving inner context
+    Callback: from inner context
+    Callback: from outer context
+    Leaving outer context
+
+As the output from the example shows, reusing a single stack object across
+multiple with statements works correctly, but attempting to nest them
+will cause the stack to be cleared at the end of the innermost with
+statement, which is unlikely to be desirable behaviour.
+
+Using separate :class:`ExitStack` instances instead of reusing a single
+instance avoids that problem::
+
+    >>> from contextlib import ExitStack
+    >>> with ExitStack() as outer_stack:
+    ...     outer_stack.callback(print, "Callback: from outer context")
+    ...     with ExitStack() as inner_stack:
+    ...         inner_stack.callback(print, "Callback: from inner context")
+    ...         print("Leaving inner context")
+    ...     print("Leaving outer context")
+    ...
+    Leaving inner context
+    Callback: from inner context
+    Leaving outer context
+    Callback: from outer context
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index fb89118..d3219f6 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -166,20 +166,16 @@ class redirect_stdout:
 
     def __init__(self, new_target):
         self._new_target = new_target
-        self._old_target = self._sentinel = object()
+        # We use a list of old targets to make this CM re-entrant
+        self._old_targets = []
 
     def __enter__(self):
-        if self._old_target is not self._sentinel:
-            raise RuntimeError("Cannot reenter {!r}".format(self))
-        self._old_target = sys.stdout
+        self._old_targets.append(sys.stdout)
         sys.stdout = self._new_target
         return self._new_target
 
     def __exit__(self, exctype, excinst, exctb):
-        restore_stdout = self._old_target
-        self._old_target = self._sentinel
-        sys.stdout = restore_stdout
-
+        sys.stdout = self._old_targets.pop()
 
 
 class suppress:
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index 916ac80..b8770c8 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -666,11 +666,18 @@ class TestRedirectStdout(unittest.TestCase):
         obj = redirect_stdout(None)
         self.assertEqual(obj.__doc__, cm_docstring)
 
+    def test_no_redirect_in_init(self):
+        orig_stdout = sys.stdout
+        redirect_stdout(None)
+        self.assertIs(sys.stdout, orig_stdout)
+
     def test_redirect_to_string_io(self):
         f = io.StringIO()
         msg = "Consider an API like help(), which prints directly to stdout"
+        orig_stdout = sys.stdout
         with redirect_stdout(f):
             print(msg)
+        self.assertIs(sys.stdout, orig_stdout)
         s = f.getvalue().strip()
         self.assertEqual(s, msg)
 
@@ -682,23 +689,26 @@ class TestRedirectStdout(unittest.TestCase):
     def test_cm_is_reusable(self):
         f = io.StringIO()
         write_to_f = redirect_stdout(f)
+        orig_stdout = sys.stdout
         with write_to_f:
             print("Hello", end=" ")
         with write_to_f:
             print("World!")
+        self.assertIs(sys.stdout, orig_stdout)
         s = f.getvalue()
         self.assertEqual(s, "Hello World!\n")
 
-    # If this is ever made reentrant, update the reusable-but-not-reentrant
-    # example at the end of the contextlib docs accordingly.
-    def test_nested_reentry_fails(self):
+    def test_cm_is_reentrant(self):
         f = io.StringIO()
         write_to_f = redirect_stdout(f)
-        with self.assertRaisesRegex(RuntimeError, "Cannot reenter"):
+        orig_stdout = sys.stdout
+        with write_to_f:
+            print("Hello", end=" ")
             with write_to_f:
-                print("Hello", end=" ")
-                with write_to_f:
-                    print("World!")
+                print("World!")
+        self.assertIs(sys.stdout, orig_stdout)
+        s = f.getvalue()
+        self.assertEqual(s, "Hello World!\n")
 
 
 class TestSuppress(unittest.TestCase):
diff --git a/Misc/NEWS b/Misc/NEWS
index c90093d..9781e1f 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -31,6 +31,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #19403: contextlib.redirect_stdout is now reentrant
+
 - Issue #19286: Directories in ``package_data`` are no longer added to
   the filelist, preventing failure outlined in the ticket.