Issue #22389: Add contextlib.redirect_stderr().

5 files changed, 85 insertions, 34 deletions

diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
index 6f36864..550b347 100644
--- a/Doc/library/contextlib.rst
+++ b/Doc/library/contextlib.rst
@@ -172,6 +172,16 @@ Functions and classes provided:
    .. versionadded:: 3.4
 
 
+.. function:: redirect_stderr(new_target)
+
+   Similar to :func:`~contextlib.redirect_stdout` but redirecting
+   :data:`sys.stderr` to another file or file-like object.
+
+   This context manager is :ref:`reentrant <reentrant-cms>`.
+
+   .. versionadded:: 3.5
+
+
 .. class:: ContextDecorator()
 
    A base class that enables a context manager to also be used as a decorator.
diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst
index 079c80f..658c975 100644
--- a/Doc/whatsnew/3.5.rst
+++ b/Doc/whatsnew/3.5.rst
@@ -148,6 +148,15 @@ compileall
   can now do parallel bytecode compilation.
   (Contributed by Claudiu Popa in :issue:`16104`.)
 
+contextlib
+----------
+
+* The new :func:`contextlib.redirect_stderr` context manager(similar to
+  :func:`contextlib.redirect_stdout`) makes it easier for utility scripts to
+  handle inflexible APIs that write their output to :data:`sys.stderr` and
+  don't provide any options to redirect it.
+  (Contributed by Berker Peksag in :issue:`22389`.)
+
 doctest
 -------
 
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index 82ee955..2fbc90c 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -5,7 +5,7 @@ from collections import deque
 from functools import wraps
 
 __all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack",
-           "redirect_stdout", "suppress"]
+           "redirect_stdout", "redirect_stderr", "suppress"]
 
 
 class ContextDecorator(object):
@@ -151,8 +151,27 @@ class closing(object):
     def __exit__(self, *exc_info):
         self.thing.close()
 
-class redirect_stdout:
-    """Context manager for temporarily redirecting stdout to another file
+
+class _RedirectStream:
+
+    _stream = None
+
+    def __init__(self, new_target):
+        self._new_target = new_target
+        # We use a list of old targets to make this CM re-entrant
+        self._old_targets = []
+
+    def __enter__(self):
+        self._old_targets.append(getattr(sys, self._stream))
+        setattr(sys, self._stream, self._new_target)
+        return self._new_target
+
+    def __exit__(self, exctype, excinst, exctb):
+        setattr(sys, self._stream, self._old_targets.pop())
+
+
+class redirect_stdout(_RedirectStream):
+    """Context manager for temporarily redirecting stdout to another file.
 
         # How to send help() to stderr
         with redirect_stdout(sys.stderr):
@@ -164,18 +183,13 @@ class redirect_stdout:
                 help(pow)
     """
 
-    def __init__(self, new_target):
-        self._new_target = new_target
-        # We use a list of old targets to make this CM re-entrant
-        self._old_targets = []
+    _stream = "stdout"
 
-    def __enter__(self):
-        self._old_targets.append(sys.stdout)
-        sys.stdout = self._new_target
-        return self._new_target
 
-    def __exit__(self, exctype, excinst, exctb):
-        sys.stdout = self._old_targets.pop()
+class redirect_stderr(_RedirectStream):
+    """Context manager for temporarily redirecting stderr to another file."""
+
+    _stream = "stderr"
 
 
 class suppress:
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index 39cc776..c52066b 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -718,60 +718,76 @@ class TestExitStack(unittest.TestCase):
         stack.push(cm)
         self.assertIs(stack._exit_callbacks[-1], cm)
 
-class TestRedirectStdout(unittest.TestCase):
+
+class TestRedirectStream:
+
+    redirect_stream = None
+    orig_stream = None
 
     @support.requires_docstrings
     def test_instance_docs(self):
         # Issue 19330: ensure context manager instances have good docstrings
-        cm_docstring = redirect_stdout.__doc__
-        obj = redirect_stdout(None)
+        cm_docstring = self.redirect_stream.__doc__
+        obj = self.redirect_stream(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)
+        orig_stdout = getattr(sys, self.orig_stream)
+        self.redirect_stream(None)
+        self.assertIs(getattr(sys, self.orig_stream), 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)
+        orig_stdout = getattr(sys, self.orig_stream)
+        with self.redirect_stream(f):
+            print(msg, file=getattr(sys, self.orig_stream))
+        self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
         s = f.getvalue().strip()
         self.assertEqual(s, msg)
 
     def test_enter_result_is_target(self):
         f = io.StringIO()
-        with redirect_stdout(f) as enter_result:
+        with self.redirect_stream(f) as enter_result:
             self.assertIs(enter_result, f)
 
     def test_cm_is_reusable(self):
         f = io.StringIO()
-        write_to_f = redirect_stdout(f)
-        orig_stdout = sys.stdout
+        write_to_f = self.redirect_stream(f)
+        orig_stdout = getattr(sys, self.orig_stream)
         with write_to_f:
-            print("Hello", end=" ")
+            print("Hello", end=" ", file=getattr(sys, self.orig_stream))
         with write_to_f:
-            print("World!")
-        self.assertIs(sys.stdout, orig_stdout)
+            print("World!", file=getattr(sys, self.orig_stream))
+        self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
         s = f.getvalue()
         self.assertEqual(s, "Hello World!\n")
 
     def test_cm_is_reentrant(self):
         f = io.StringIO()
-        write_to_f = redirect_stdout(f)
-        orig_stdout = sys.stdout
+        write_to_f = self.redirect_stream(f)
+        orig_stdout = getattr(sys, self.orig_stream)
         with write_to_f:
-            print("Hello", end=" ")
+            print("Hello", end=" ", file=getattr(sys, self.orig_stream))
             with write_to_f:
-                print("World!")
-        self.assertIs(sys.stdout, orig_stdout)
+                print("World!", file=getattr(sys, self.orig_stream))
+        self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
         s = f.getvalue()
         self.assertEqual(s, "Hello World!\n")
 
 
+class TestRedirectStdout(TestRedirectStream, unittest.TestCase):
+
+    redirect_stream = redirect_stdout
+    orig_stream = "stdout"
+
+
+class TestRedirectStderr(TestRedirectStream, unittest.TestCase):
+
+    redirect_stream = redirect_stderr
+    orig_stream = "stderr"
+
+
 class TestSuppress(unittest.TestCase):
 
     @support.requires_docstrings
diff --git a/Misc/NEWS b/Misc/NEWS
index 1cf01c9..86762f8 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -191,6 +191,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #22389: Add contextlib.redirect_stderr().
+
 - Issue #21356: Make ssl.RAND_egd() optional to support LibreSSL. The
   availability of the function is checked during the compilation. Patch written
   by Bernard Spil.