Is it possible to access enclosing context manager?

Viewed 5052

There are essentially three ways to use the with statement:

Use an existing context manager:

with manager:
    pass

Create a context manager and bind its result to a variable:

with Manager() as result:
    pass

Create an context manager and discard its return value:

with Manager():
    pass

If we have place a function get_manager() inside the three with blocks above, is there any implementation that can return the enclosing context manager, or at least their __exit__ function?

It's obviously easy in the first case, but I can't think of a way to make it work in the other two. I doubt it's possible to get the entire context manager, since the value stack is popped immediately after the SETUP_WITH opcode. However, since the __exit__ function is stored on the block stack by SETUP_WITH, is there some way to access it?

5 Answers

If you will accept a hacky solution, I bring you one inspired by this.

Have the context manager edit the local namespace.

class Context(object):
    def __init__(self, locals_reference):
        self.prev_context = locals_reference.get('context', None)
        self.locals_reference = locals_reference

    def __enter__(self):
        self.locals_reference['context'] = self

    def __exit__(self, exception_type, exception_value, traceback):
        if self.prev_context is not None:
            self.locals_reference['context'] = self.prev_context
        else:
            del self.locals_reference['context']

You can then get the context with the context variable

with Context(locals()):
    print(context)

This implementation also works on nested contexts

with Context(locals()):
    c_context = context
    with Context(locals()):
        print(c_context == context)
    print(c_context == context)

However this is implementation specific, as the return value of locals may be a copy of the namespace. Tested on CPython 3.10.

Edit:

The implementation above will not work in functions from other modules (I wonder why), so here is a function that fetches the context:

def get_current_context(cls) -> "Context | None":
    try:
        if context is not None:
            return context
    except NameError:
        pass

    i = 0
    while True:
        try:
            c = sys._getframe(i).f_locals.get('context',None)
        except ValueError:
            return None

        if c is not None:
            return c
        i += 1

I would make it a classmethod of the context manager class.

Many years late, here's a straightforward way to do this in the second of the OP's 3 cases, where with ... as is used to bind the output of the context manager's __enter__ method to a variable. You can have the __enter__ method return the context manager itself, or its __exit__ method if that's all you're interested in.

class Manager:
     def __enter__(self): return self
     def __exit__(self, *args): print("Doing exit method stuff!")

with Manager() as manager:
    print("Doing some stuff before manually calling the exit method...")
    manager.__exit__() # call the exit method
    print("Doing some more stuff before exiting for real...")

Of course, this would interfere with using with ... as to bind some other return-value from __enter__, but it would be straightforward to have __enter__ return a tuple consisting of its ordinary return value and the manager, just as you can make a function return multiple values.

As the OP noted, it's also straightforward to call the __exit__ method in the first case, where the context manager had already been assigned to a variable beforehand. So the only really tricky case is the third one where the context manager is simply created via with Manager(): but is never assigned to a variable. My advice would be: if you're going to want to refer to the manager (or its methods) later, then either (1) assign it a name beforehand, (2) have its __enter__ method return a reference to it for with ... as to capture as I did above, but (3) DO NOT create it without storing any reference to it!

Related