Welcome, guest | Sign In | My Account | Store | Cart

Use the "with" keyword or a decorator to simplify a bit redirecting IO to a file.

Python, 56 lines
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
#!/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

If no file or StringIO (or any file-like) is given, it will redirect the output to a StringIO(). The default opening mode (a+) is not arbitrary: you will probably want to append each printed line to the file and still be able to read it. After the "with" block, the file is closed (by default).

3 comments

Paul Moore 14 years ago  # | flag

You should restore the old stdout. Rather than setting sys.stdout to sys.__stdout__ in __exit__, you should save the old sys.stdout in an instance variable in __enter__, and restore it in __exit__. That way, if the user nests this construct, or redirects sys.stdout some other way, you won't clobber his change.

Nice point, thanks!

Paul Moore 14 years ago  # | flag

Typo in the __exit__ code: sys.stdout = sys.original_stdout. That should be self.original_stdout.