Capture stdout and stderr conveniently from within Python programs

gistfile1.py

import contextlib
import sys
from StringIO import StringIO

# Output capturing helpers thanks to
# test.test_support and http://stackoverflow.com/a/18844643/240490
# extended by Jonathan Eunice to add easy saving to a file (including
# auto-save if a filepath is provided in the constructor)

# This is not quite polished enough to make it into a full PyPI 
# project, but more convenient and developed an idea than should
# be cast aside. Thus, a Gist.

class SaveableStringIO(StringIO):
    def __init__(self, filepath=None, filemode="w"):
        StringIO.__init__(self)
        self.filepath = filepath
        self.filemode = filemode

    def save(self, filepath=None, filemode=None):
        """
        Easy write-to-file capability for a StringIO.
        The filepath and filemode may be given directly,
        or grabbed from saved values from when the instance
        was created. At least one source of filepath is
        required.
        """
        filepath = filepath or self.filepath
        if filepath is None:
            raise IOError("must set or give filepath")
        filemode = filemode or self.filemode
        with open(filepath, filemode) as f:
            f.write(self.getvalue())

@contextlib.contextmanager
def captured_output(stream_name, filepath=None, filemode="w"):
    """
    Return a context manager used by captured_stdout and captured_stdin
    that temporarily replaces the sys stream *stream_name* with a
    SaveableStringIO. If a filepath is given, output is written there
    directly.
    """
    orig_stdout = getattr(sys, stream_name)
    sio = SaveableStringIO(filepath, filemode)
    setattr(sys, stream_name, sio)
    try:
        yield getattr(sys, stream_name)
    finally:
        if sio.filepath is not None:
            sio.save()
        setattr(sys, stream_name, orig_stdout)

def captured_stdout(filepath=None, filemode="w"):
    """Capture the output of sys.stdout:

       with captured_stdout() as s:
           print "hello"
       self.assertEqual(s.getvalue(), "hello")
    """
    return captured_output("stdout", filepath)

def captured_stderr(filepath=None, filemode="w"):
    return captured_output("stderr", filepath, filemode)

def captured_stdin(filepath=None, filemode="w"):
    return captured_output("stdin", filepath, filemode)

# Might want to hande the case where stream_name is a StringIO,
# so multiple captures can be chained, like iocapture on PyPI