class CaptureStdio(object):
"""
A simple context manager for redirecting stdout and stderr to a StringIO
buffer.
"""
def __init__(self, stdout=True, stderr=True):
self.stdout = StringIO()
self.stderr = StringIO()
def __enter__(self):
self._original_stdout = sys.stdout
self._original_stderr = sys.stderr
sys.stdout = self.stdout
sys.stderr = self.stderr
return self.stdout, self.stderr
def __exit__(self, *args, **kwargs):
sys.stdout = self._original_stdout
sys.stderr = self._original_stderr
self.stdout.close()
self.stderr.close()
