ActiveState Code

#!/usr/bin/python
from __future__ import with_statement

import sys
from StringIO import StringIO

__all__ = ['RedirectedIO', 'redirect_io']


class RedirectedIO(object):
    def __init__(self, target=None, mode='a+',
                 close_target=True):
        try:
            target = open(target, mode)
        except TypeError:
            if target is None:
                target = StringIO()
        self.target = target
        self.close_target = close_target

    def __enter__(self):
        """ Redirect IO to self.target.
        """
        self.original_stdout = sys.stdout
        sys.stdout = self.target
        return self.target

    def __exit__(self, *args, **kwargs):
        """ Restore stdio and close the file.
        """
        sys.stdout = self.original_stdout
        if self.close_target:
            self.target.close()


def redirect_io(target=None, mode='a+', keep_target=True):
    """ Returns a decorator that wrapps a
    function and redirects its IO to a target
    file (a StringIO by default). The target is
    available as .iotarget on the decorated function.
    """
    def dec(func):

        def wrapper(*args, **kwargs):
            with RedirectedIO(target, mode, not keep_target) as iotarget:
                result = func(*args, **kwargs)
                if keep_target:
                    wrapper.iotarget = iotarget
            return result

        wrapper.iotarget = None
        wrapper.__doc__ = func.__doc__
        wrapper.__name__ = func.__name__
        return wrapper

    return dec