Temporary substitution of object in modules via with statement (Python recipe) by Jacob Oscarson

ActiveState Code (http://code.activestate.com/recipes/577020/)

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()