capture_output.py

class CaptureOutput(object):
    def __init__(self, *what):
        self.__names = set(what)
        self.__objects = {}
        for name in self.__names:
            assert isinstance(name, str)
            assert name in ('stderr', 'stdout', '__stderr__', '__stdout__')

    def __enter__(self):
        for name in self.__names:
            self.__objects[name] = getattr(sys, name)
            buff = StringIO()
            setattr(sys, name, buff)
            setattr(self, name, buff)
        return self

    def __exit__(self, evalue=None, etype=None, tb=None):
        for name, original in self.__objects.items():
            setattr(sys, name, original)
            setattr(self, name, getattr(self, name).getvalue())