130

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
Improve this question
0

12 Answers 12

Reset to default
101

On Python 3.4+, use the contextlib.redirect_stdout context manager:

from contextlib import redirect_stdout import io f = io.StringIO() with redirect_stdout(f): help(pow) s = f.getvalue()
Improve this answer
Sign up to request clarification or add additional context in comments.

1 Comment

This does not solve the problem when trying to write to sys.stdout.buffer (as you need to do when writing bytes). StringIO does not have the buffer attribute, while TextIOWrapper does. See the answer from @JonnyJD.
58

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()
Improve this answer

2 Comments

check_output directly captures the output of a command run in a subprocess: <br> value = subprocess.check_output( command, shell=True)
On modern Python versions you can do capture_output=True instead of stdout=subprocess.PIPE.
52

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'
Improve this answer

5 Comments

Love this solution. I modified, so as not to accidentally lose stuff from a stream on which I'm not expecting output, e.g. unexpected errors. In my case, capture() can accept sys.stderr or sys.stdout as a parameter, indicating to only capture that stream.
StringIO doesn't support unicode in any fashion, so you can integrate the answer here to make the above support non-ASCII chars: stackoverflow.com/a/1819009/425050
Modifying a yielded value in the finally is really rather wierd - with capture() as out: will behave differently to with capture() as out, err:
Unicode / stdout.buffer support can be reached with using the io module. See my answer.
This solution breaks if you use subprocess and redirect output to sys.stdout/stderr. This is because StringIO is not a real file object and misses the fileno() function.
23
+100

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.

Improve this answer

Comments

14

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
Improve this answer

Comments

10

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
Improve this answer

Comments

7

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
Improve this answer

Comments

5

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.

Improve this answer

Comments

4

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")
Improve this answer

Comments

4

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
Improve this answer

Comments

0

Another way when third party code has already copied a reference to sys.stdout is to temporarily replace the write() method itself:

from types import MethodType ... f = io.StringIO() def new_write(self, data): f.write(data) old_write = sys.stdout.write sys.stdout.write = MethodType(new_write, sys.stdout) error = command.run(args) sys.stdout.write = old_write output = f.getvalue()
Improve this answer

Comments

0

Thought that @arthur's comment should exist in an answer.
Using the check_output method from subprocess seems easiest:

In [1]: import subprocess ...: ...: command = "echo 'hello world'" ...: output = subprocess.check_output(command, shell=True, encoding='utf-8') In [2]: print(output) hello world
Improve this answer

Comments

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.