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