Issue #15805: Add contextlib.redirect_stdout()

rhettinger · rhettinger · commit 088cbf2d390f · 2013-10-10T00:46:57.000-07:00
diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
@@ -115,6 +115,37 @@ Functions and classes provided:
 
    .. versionadded:: 3.4
 
+.. function:: redirect_stdout(new_target)
+
+   Context manager for temporarily redirecting :data:`sys.stdout` to
+   another file or file-like object.
+
+   This tool adds flexibility to existing functions or classes whose output
+   is hardwired to stdout.
+
+   For example, the output of :func:`help` normally is sent to *sys.stdout*.
+   You can capture that output in a string by redirecting the output to a
+   :class:`io.StringIO` object::
+
+        f = io.StringIO()
+        with redirect_stdout(f):
+            help(pow)
+        s = f.getvalue()
+
+   To send the output of :func:`help` to a file on disk, redirect the output
+   to a regular file::
+
+        with open('help.txt', 'w') as f:
+            with redirect_stdout(f):
+                help(pow)
+
+   To send the output of :func:`help` to *sys.stderr*::
+
+        with redirect_stdout(sys.stderr):
+            help(pow)
+
+   .. versionadded:: 3.4
+
 .. class:: ContextDecorator()
 
    A base class that enables a context manager to also be used as a decorator.
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
@@ -4,7 +4,8 @@
 from collections import deque
 from functools import wraps
 
-__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack", "ignored"]
+__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack",
+           "ignored", "redirect_stdout"]
 
 
 class ContextDecorator(object):
@@ -140,6 +141,43 @@ def __enter__(self):
     def __exit__(self, *exc_info):
         self.thing.close()
 
+class redirect_stdout:
+    """Context manager for temporarily redirecting stdout to another file
+
+        # How to send help() to stderr
+
+        with redirect_stdout(sys.stderr):
+            help(dir)
+
+        # How to write help() to a file
+
+        with open('help.txt', 'w') as f:
+            with redirect_stdout(f):
+                help(pow)
+
+        # How to capture disassembly to a string
+
+        import dis
+        import io
+
+        f = io.StringIO()
+        with redirect_stdout(f):
+            dis.dis('x**2 - y**2')
+        s = f.getvalue()
+
+    """
+
+    def __init__(self, new_target):
+        self.new_target = new_target
+
+    def __enter__(self):
+        self.old_target = sys.stdout
+        sys.stdout = self.new_target
+        return self.new_target
+
+    def __exit__(self, exctype, excinst, exctb):
+        sys.stdout = self.old_target
+
 @contextmanager
 def ignored(*exceptions):
     """Context manager to ignore specified exceptions
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
@@ -1,5 +1,6 @@
 """Unit tests for contextlib.py, and other context managers."""
 
+import io
 import sys
 import tempfile
 import unittest
@@ -653,6 +654,14 @@ def test_exception_hierarchy(self):
         with ignored(LookupError):
             'Hello'[50]
 
+class TestRedirectStdout(unittest.TestCase):
+
+    def test_redirect_to_string_io(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            help(pow)
+        s = f.getvalue()
+        self.assertIn('pow', s)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -32,6 +32,8 @@ Library
 - Issue #19158:  a rare race in BoundedSemaphore could allow .release() too
   often.
 
+- Issue #15805: Add contextlib.redirect_stdout().
+
 - Issue #18716: Deprecate the formatter module.
 
 - Issue #18037: 2to3 now escapes '\u' and '\U' in native strings.
