Question or problem about Python programming:
suppose there is a script doing something like this:
# module writer.py
import sys
def write():
sys.stdout.write("foobar")
Now suppose I want to capture the output of the write function and store it in a variable for further processing. The naive solution was:
# module mymodule.py from writer import write out = write() print out.upper()
But this doesn’t work. I come up with another solution and it works, but please, let me know if there is a better way to solve the problem. Thanks
import sys from cStringIO import StringIO # setup the environment backup = sys.stdout # #### sys.stdout = StringIO() # capture output write() out = sys.stdout.getvalue() # release output # #### sys.stdout.close() # close the stream sys.stdout = backup # restore original stdout print out.upper() # post processing
How to solve the problem:
Solution 1:
Setting stdout is a reasonable way to do it. Another is to run it as another process:
import subprocess proc = subprocess.Popen(["python", "-c", "import writer; writer.write()"], stdout=subprocess.PIPE) out = proc.communicate()[0] print out.upper()
Solution 2:
Here is a context manager version of your code. It yields a list of two values; the first is stdout, the second is stderr.
import contextlib
@contextlib.contextmanager
def capture():
import sys
from cStringIO import StringIO
oldout,olderr = sys.stdout, sys.stderr
try:
out=[StringIO(), StringIO()]
sys.stdout,sys.stderr = out
yield out
finally:
sys.stdout,sys.stderr = oldout, olderr
out[0] = out[0].getvalue()
out[1] = out[1].getvalue()
with capture() as out:
print 'hi'
Solution 3:
For future visitors: Python 3.4 contextlib provides for this directly (see Python contextlib help) via the redirect_stdout context manager:
from contextlib import redirect_stdout
import io
f = io.StringIO()
with redirect_stdout(f):
help(pow)
s = f.getvalue()
Solution 4:
This is the decorator counterpart of my original code.
writer.py remains the same:
import sys
def write():
sys.stdout.write("foobar")
mymodule.py sligthly gets modified:
from writer import write as _write
from decorators import capture
@capture
def write():
return _write()
out = write()
# out post processing...
And here is the decorator:
def capture(f):
"""
Decorator to capture standard output
"""
def captured(*args, **kwargs):
import sys
from cStringIO import StringIO
# setup the environment
backup = sys.stdout
try:
sys.stdout = StringIO() # capture output
f(*args, **kwargs)
out = sys.stdout.getvalue() # release output
finally:
sys.stdout.close() # close the stream
sys.stdout = backup # restore original stdout
return out # captured output wrapped in a string
return captured
Solution 5:
Or maybe use functionality that is already there…
from IPython.utils.capture import capture_output
with capture_output() as c:
print('some output')
c()
print c.stdout
Solution 6:
Starting with Python 3 you can also use sys.stdout.buffer.write() to write (already) encoded byte strings to stdout (see stdout in Python 3).
When you do that, the simple StringIO approach doesn’t work because neither sys.stdout.encoding nor sys.stdout.buffer would be available.
Starting with Python 2.6 you can use the TextIOBase API, which includes the missing attributes:
import sys
from io import TextIOWrapper, BytesIO
# setup the environment
old_stdout = sys.stdout
sys.stdout = TextIOWrapper(BytesIO(), sys.stdout.encoding)
# do some writing (indirectly)
write("blub")
# get output
sys.stdout.seek(0) # jump to the start
out = sys.stdout.read() # read output
# restore stdout
sys.stdout.close()
sys.stdout = old_stdout
# do stuff with the output
print(out.upper())
This solution works for Python 2 >= 2.6 and Python 3.
Please note that our sys.stdout.write() only accepts unicode strings and sys.stdout.buffer.write() only accepts byte strings.
This might not be the case for old code, but is often the case for code that is built to run on Python 2 and 3 without changes.
If you need to support code that sends byte strings to stdout directly without using stdout.buffer, you can use this variation:
class StdoutBuffer(TextIOWrapper):
def write(self, string):
try:
return super(StdoutBuffer, self).write(string)
except TypeError:
# redirect encoded byte strings directly to buffer
return super(StdoutBuffer, self).buffer.write(string)
You don’t have to set the encoding of the buffer the sys.stdout.encoding, but this helps when using this method for testing/comparing script output.
Solution 7:
The question here (the example of how to redirect output, not the tee part) uses os.dup2 to redirect a stream at the OS level. That is nice because it will apply to commands that you spawn from your program as well.
Solution 8:
I think You should look at these four objects:
from test.test_support import captured_stdout, captured_output, \
captured_stderr, captured_stdin
Example:
from writer import write
with captured_stdout() as stdout:
write()
print stdout.getvalue().upper()
UPD: As Eric said in a comment, one shouldn’t use they directly, so I copied and pasted it.
# Code from test.test_support:
import contextlib
import sys
@contextlib.contextmanager
def captured_output(stream_name):
"""Return a context manager used by captured_stdout and captured_stdin
that temporarily replaces the sys stream *stream_name* with a StringIO."""
import StringIO
orig_stdout = getattr(sys, stream_name)
setattr(sys, stream_name, StringIO.StringIO())
try:
yield getattr(sys, stream_name)
finally:
setattr(sys, stream_name, orig_stdout)
def captured_stdout():
"""Capture the output of sys.stdout:
with captured_stdout() as s:
print "hello"
self.assertEqual(s.getvalue(), "hello")
"""
return captured_output("stdout")
def captured_stderr():
return captured_output("stderr")
def captured_stdin():
return captured_output("stdin")
Solution 9:
I like the contextmanager solution however if you need the buffer stored with the open file and fileno support you could do something like this.
import six
from six.moves import StringIO
class FileWriteStore(object):
def __init__(self, file_):
self.__file__ = file_
self.__buff__ = StringIO()
def __getattribute__(self, name):
if name in {
"write", "writelines", "get_file_value", "__file__",
"__buff__"}:
return super(FileWriteStore, self).__getattribute__(name)
return self.__file__.__getattribute__(name)
def write(self, text):
if isinstance(text, six.string_types):
try:
self.__buff__.write(text)
except:
pass
self.__file__.write(text)
def writelines(self, lines):
try:
self.__buff__.writelines(lines)
except:
pass
self.__file__.writelines(lines)
def get_file_value(self):
return self.__buff__.getvalue()
use
import sys sys.stdout = FileWriteStore(sys.stdout) print "test" buffer = sys.stdout.get_file_value() # you don't want to print the buffer while still storing # else it will double in size every print sys.stdout = sys.stdout.__file__ print buffer
Solution 10:
Here’s a context manager taking inspiration from @JonnyJD’s answer supporting writing bytes to buffer attributes abut also taking advantage of sys’s dunder-io referenes for further simplification.
import io
import sys
import contextlib
@contextlib.contextmanager
def capture_output():
output = {}
try:
# Redirect
sys.stdout = io.TextIOWrapper(io.BytesIO(), sys.stdout.encoding)
sys.stderr = io.TextIOWrapper(io.BytesIO(), sys.stderr.encoding)
yield output
finally:
# Read
sys.stdout.seek(0)
sys.stderr.seek(0)
output['stdout'] = sys.stdout.read()
output['stderr'] = sys.stderr.read()
sys.stdout.close()
sys.stderr.close()
# Restore
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__
with capture_output() as output:
print('foo')
sys.stderr.buffer.write(b'bar')
print('stdout: {stdout}'.format(stdout=output['stdout']))
print('stderr: {stderr}'.format(stderr=output['stderr']))
Output is:
stdout: foo stderr: bar
