用于临时变量赋值的 Python 上下文管理器

Question

I'm often in the need of temporarily switching the value of a variable out with something else, do some computation which depends on this variable, and then restore the variable to its original value. Eg:

var = 0
# Assign temporary value and do computation
var_ori = var
var = 1
do_something_with_var()  # Function that reads the module level var variable
# Reassign original value
var = var_ori

This seems like an obvious opportunity for using a context manager (the with statement). Does the Python standard library contain any such context manager?

Edit

I am aware that this sort of thing is most often handled by other, better means than temporarily changing a variable. I am however not asking for obvious workarounds.

In my actual work case, I can not alter the do_something_with_var function. Actually this is not even a function, but a string of code which gets evaluated in the context of the global namespace as part of some metaprogramming. The example I gave was the simplest I could think of that kept my problem with the temporary variable. I did not ask to get a workaround (proper version) of my example code, but rather to get an answer on my written question.

Answer 1

Nope, because a context manager can't assign variables in the caller's scope like that. (Anyone who thinks you could do it with locals or inspect , try using the context manager you come up with inside a function. It won't work.)

There are utilities for doing that with things that aren't local variables, such as module globals, other object attributes, and dicts... but they're unittest.mock.patch and its related functions, so you should strongly consider other alternatives before using them in a non-testing context. Operations like "temporarily modify this thing and then restore it" tend to lead to confusing code, and may indicate you're using too much global state.

Answer 2

The simple answer to your question:

Does the Python standard library contain any such context manager?

is "No, it does not."

Answer 3

My mistake, instead perhaps something like this, it is not built-in:

class ContextTester(object):
    """Initialize context environment and replace variables when completed"""

    def __init__(self, locals_reference):
        self.prev_local_variables = locals_reference.copy()
        self.locals_reference = locals_reference

    def __enter__(self):
        pass

    def __exit__(self, exception_type, exception_value, traceback):
        self.locals_reference.update(self.prev_local_variables)



a = 5
def do_some_work():
    global a
    print(a)
    a = 8
print("Before context tester: {}".format(a))
with ContextTester(locals()) as context:
    print("In context tester before assignment: {}".format(a))
    a = 6
    do_some_work()
    print("In context tester after assignment: {}".format(a))
print("After context tester: {}".format(a))

Output:

Before context tester: 5
In context tester before assignment: 5
6
In context tester after assignment: 8
After context tester: 5

For clarity, so you know it's actually doing something:

class ContextTester(object):
    """Initialize context environment and replace variables when completed"""

    def __init__(self, locals_reference):
        self.prev_local_variables = locals_reference.copy()
        self.locals_reference = locals_reference

    def __enter__(self):
        pass

    def __exit__(self, exception_type, exception_value, traceback):
        #self.locals_reference.update(self.prev_local_variables)
        pass

a = 5
def do_some_work():
    global a
    print(a)
    a = 8
print("Before context tester: {}".format(a))
with ContextTester(locals()) as context:
    print("In context tester before assignment: {}".format(a))
    a = 6
    do_some_work()
    print("In context tester after assignment: {}".format(a))
print("After context tester: {}".format(a))
a = 5
print("Before context tester: {}".format(a))
with ContextTester(locals()) as context:
    print("In context tester before assignment: {}".format(a))
    a = 6
    print("In context tester after assignment: {}".format(a))
print("After context tester: {}".format(a))

Output:

Before context tester: 5
In context tester before assignment: 5
6
In context tester after assignment: 8
After context tester: 8

Before context tester: 5
In context tester before assignment: 5
In context tester after assignment: 6
After context tester: 6

You could also do this:

def wrapper_function(func, *args, **kwargs):
    prev_globals = globals().copy()
    func(*args, **kwargs)
    globals().update(prev_globals)

It should be noted that if you try to use the with statement within a function, you'll want to use globals() as the reference to locals and it may have unintended consequences, might still anyways.

I wouldn't recommend doing this at all, but should work.