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

See docstring in the code

Python, 67 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
57
58
59
60
61
62
63
64
65
66
67
from __future__ import with_statement # only needed on python 2.5

import sys
from functools import partial

def _findattr(mod, rest):
    parent, dot, rest = rest.partition('.')
    if rest:
        return _findattr(getattr(mod, parent), rest)
    else:
        return mod, parent

def find_in_module(fullname):
    modname, dot, rest = fullname.partition('.')
    module, objname = _findattr(__import__(modname), rest)
    return getattr(module, objname), partial(setattr, module, objname)


class substitute(object):
    """
    A context manager that takes the name of a globally reachable
    object (in the form of 'module.object', e.g. `sys.stdout`) and
    substitutes it with another object while the context manager is in
    effect.

    Example::

        >>> from StringIO import StringIO
        >>> capture = StringIO()
        >>> with substitute('sys.stdout', capture):
        ...     print('foo')
        >>> capture.getvalue()
        'foo\\n'

    or::

        >>> import os
        >>> with substitute('os.path.exists', lambda p: 'Yes indeedy!'):
        ...     assert os.path.exists('/no/such/path') == 'Yes indeedy!'
        >>> assert os.path.exists('/no/such/path') == False

    Exceptions are propagated after the value is restored::

        >>> import os
        >>> with substitute('os.environ', {}):
        ...    os.environ['PATH']
        Traceback (most recent call last):
        KeyError
    """

    def __init__(self, name, substitution):
        self.name = name
        self.substitution = substitution

    def __enter__(self):
        self.oldvalue, self._set = find_in_module(self.name)
        self._set(self.substitution)
        return self

    def __exit__(self, exc, value, tb):
        self._set(self.oldvalue)
        if tb is not None:
            raise(exc, value, tb)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

At my work (Open End), we use a method of substituting stuff in a global namespace of our Javascript-heavy application. When working on a personal project I found myself wanting to do the same in Python. This is how I implemented a similar functionality using the with statement.

It's been quite handy for me, especially in combination with Voidspace's mock library.

As of now, it doesn't support replacing a whole module, like this:

>>> with substitute('os', my_faked_os_mod):
       assert os.linesep == 'my replaced linesep'

That might be possible employing some of the tricks found in Ryan Kelly's withhacks, but I thought that a little bit to much, at least for something that's a 'recipie'.

Improvements, anyone?

1 comment

George Sakkis 14 years, 2 months ago  # | flag