Issue #15805: Add contextlib.redirect_stdout()

2 files changed, 48 insertions, 1 deletions

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index aaab095..868fa6c 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -4,7 +4,8 @@ import sys
 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 @@ class closing(object):
     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
index 2892917..d8a0530 100644
--- 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 @@ class TestIgnored(unittest.TestCase):
         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()