Search code examples
pythonstatecontextmanagerstateful

Is it possible to make a contextually-sensitive python context manager that saves, modifies and restores state?


I have a pair of python functions that currently flip a global variable between two values. I would like to turn them into context managers so I can use them as with blocks, setting the variable inside the block, but restoring it after. Here's the desired behaviour:

>>> MODE
'user'
>>> mode_sudo()  # Sets MODE to 'sudo'...
>>> MODE
'sudo'
>>> mode_user()  # Sets MODE to 'user'...
>>> MODE
'user'
>>> with mode_sudo():
...    print MODE
'sudo'
>>> MODE
'user'

Is such a chimera possible?

UPDATE: Just for clarity, here's the context-manager-only implementation:

from contextlib import contextmanager

@contextmanager
def mode_sudo():
    global MODE
    old_mode = MODE
    MODE = 'sudo'
    yield
    MODE = old_mode

@contextmanager
def mode_user():
    global MODE
    old_mode = MODE
    MODE = 'user'
    yield
    MODE = old_mode

Calling these w/o a with keyword returns a generator. Is there a way to get the mode-flipping behavior with both the plain-vanilla function call and the chocolate context manager?


Solution

  • Do it like this:

    class mod_user:
    
        def __init__(self):
            global MODE
            self._old_mode = MODE
            MODE = "user"
    
        def __enter__(self):
            pass
    
        def __exit__(self, *args, **kws):
            global MODE
            MODE = self._old_mode
    
    MODE = "sudo"
    
    with mod_user():
        print MODE  # print : user.
    
    print MODE  # print: sudo.
    
    mod_user()
    print MODE   # print: user.