Capture stdout from a script?

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

Solution 1:

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 2:

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 3:

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'