from __future__ import with_statement from contextlib import contextmanager import sys __docformat__ = "restructuredtext" @contextmanager def restoring(expr, clone=None): '''A context manager that evaluates an expression when entering the runtime context and restores its value when exiting. This context manager makes .. python:: with restoring(expr, clone) as value: BODY a shortcut for .. python:: value = EXPR __cloned = clone(value) if clone is not None else value try: BODY finally: EXPR = __cloned del __cloned where ``__cloned`` is a temporary hidden name and ``EXPR`` is ``expr`` substituted textually in the code snippet. Therefore ``expr`` can only be an assignable expression, i.e. an expression that is allowed on the left hand side of '=' (e.g. identifier, subscription, attribute reference, etc.). :param expr: The expression whose value is to be evaluated and restored. :type expr: str :param clone: A callable that takes the object ``expr`` evaluates to and returns an appropriate copy to be used for restoring. If None, the original object is used. :type clone: callable or None ''' f = sys._getframe(2) # bypass the contextmanager frame # evaluate the expression and make a clone of the value to be restored value = eval(expr, f.f_globals, f.f_locals) restored_value = clone(value) if clone is not None else value try: yield value finally: if expr in f.f_locals: # local or nonlocal name _update_locals(f, {expr:restored_value}) elif expr in f.f_globals: # global name f.f_globals[expr] = restored_value else: # make a copy of f_locals and bind restored_value to a new name tmp_locals = dict(f.f_locals) tmp_name = '__' + min(tmp_locals) tmp_locals[tmp_name] = restored_value exec '%s = %s' % (expr, tmp_name) in f.f_globals, tmp_locals def _update_locals(frame, new_locals, clear=False): # XXX: obscure, most likely implementation-dependent fact: # f_locals can be modified (only?) from within a trace function f_trace = frame.f_trace try: sys_trace = sys.gettrace() except AttributeError: # sys.gettrace() not available before 2.6 sys_trace = None def update_tracer(frm, event, arg): # Update the frame's locals and restore both the local and the system's #trace function assert frm is frame if clear: frm.f_locals.clear() frm.f_locals.update(new_locals) frm.f_trace = f_trace sys.settrace(sys_trace) # Force tracing on with setting the global tracing function and set # the frame's local trace function sys.settrace(lambda frame, event, arg: None) frame.f_trace = update_tracer def test_restoring_immutable(): x = 'b' foo = {'a':3, 'b':4} with restoring('foo[x]') as y: assert y == foo[x] == 4 foo[x] = y = None assert y == foo[x] == None assert foo[x] == 4 and y == None assert sorted(locals()) == ['foo', 'x', 'y'] def test_restoring_mutable(): orig_path = sys.path[:] with restoring('sys.path', clone=list) as path: assert path is sys.path path += ['foo'] assert path == orig_path + ['foo'] assert sys.path == orig_path assert path == orig_path + ['foo'] assert sorted(locals()) == ['orig_path', 'path'] x = 1 def test_restoring_global(): global y; y = 2 global x with restoring('x'): x = None with restoring('y'): y += 3 assert x == None and y == 5 assert y == 2 assert x == 1 assert not locals() def test_restoring_local(): x = 5 with restoring('x'): x = None assert x == 5 assert sorted(locals()) == ['x'] def test_restoring_nonlocal(): a = [] def nested(): with restoring('a', list): a.append(1) assert a == [1] assert a == [] nested() assert a == [] if __name__ == '__main__': test_restoring_immutable() test_restoring_mutable() test_restoring_global() test_restoring_local() test_restoring_nonlocal()