For python3.4+, there's a context manager for this in the standard library.
with contextlib.redirect_stdout(file_like_object):
...
This part of the answer was updated, but is mostly for people who are still stuck in a python2.x world
If you're stuck on an older version of python, this context manager isn't too hard to write yourself. The key is that you can update sys.stdout
to whatever file-like object you want (that's what print
writes to):
>>> import sys
>>> import StringIO
>>> stdout = sys.stdout # keep a handle on the real standard output
>>> sys.stdout = StringIO.StringIO() # Choose a file-like object to write to
>>> foo()
>>> sys.stdout = stdout
>>> foo()
bar
To create a context manager to set stdout to whatever you want when you enter the context and then have the context manager reset stdout when you __exit__
the context.
Here's a simple example using contextlib
to create the context manager:
import contextlib
import sys
@contextlib.contextmanager
def stdout_redirect(where):
sys.stdout = where
try:
yield where
finally:
sys.stdout = sys.__stdout__
def foo():
print 'bar'
# Examples with StringIO
import StringIO
with stdout_redirect(StringIO.StringIO()) as new_stdout:
foo()
new_stdout.seek(0)
print "data from new_stdout:",new_stdout.read()
new_stdout1 = StringIO.StringIO()
with stdout_redirect(new_stdout1):
foo()
new_stdout1.seek(0)
print "data from new_stdout1:",new_stdout1.read()
# Now with a file object:
with open('new_stdout') as f:
with stdout_redirect(f):
foo()
# Just to prove that we actually did put stdout back as we were supposed to
print "Now calling foo without context"
foo()
Note:
On python3.x, StringIO.StringIO
has moved to io.StringIO
. Also, on python2.x, cStringIO.StringIO
might be slightly more performant.