Is it possible to access enclosing context manager?

Question

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?

Answer 1

Unfortunately, as discussed in the comments, this is not possible in all cases. When a context manager is created, the following code is run (in cPython 2.7, at least. I can't comment on other implementations):

    case SETUP_WITH:
    {
        static PyObject *exit, *enter;
        w = TOP();
        x = special_lookup(w, "__exit__", &exit);
        if (!x)
            break;
        SET_TOP(x);
        /* more code follows... */
    }

The __exit__ method is pushed onto a stack with the SET_TOP macro, which is defined as:

#define SET_TOP(v)        (stack_pointer[-1] = (v))

The stack pointer, in turn, is set to the top of the frame's value stack at the start of frame eval:

stack_pointer = f->f_stacktop;

Where f is a frame object defined in frameobject.h. Unfortunately for us, this is where the trail stops. The python accessible frame object is defined with the following methods only:

static PyMemberDef frame_memberlist[] = {
    {"f_back",          T_OBJECT,       OFF(f_back),    RO},
    {"f_code",          T_OBJECT,       OFF(f_code),    RO},
    {"f_builtins",      T_OBJECT,       OFF(f_builtins),RO},
    {"f_globals",       T_OBJECT,       OFF(f_globals), RO},
    {"f_lasti",         T_INT,          OFF(f_lasti),   RO},
    {NULL}      /* Sentinel */
};

Which, unfortunaltey, does not include the f_valuestack that we would need. This makes sense, since f_valuestack is of the type PyObject **, which would need to be wrapped in an object to be accessible from python any way.

TL;DR: The __exit__ method we're looking for is only located in one place, the value stack of a frame object, and cPython doesn't make the value stack accessible to python code.

Answer 2

The difference between this case and similar-appearing cases like super is that here there is no enclosing frame to look at. A with statement is not a new scope. sys._getframe(0) (or, if you're putting the code into a function, sys._getframe(1)) will work just fine, but it'll return you the exact same frame you have before and after the with statement.

The only way you could do it would be by inspecting the bytecode. But even that won't help. For example, try this:

from contextlib import contextmanager

@contextmanager
def silly():
    yield

with silly():
    fr = sys._getframe(0)

dis.dis(fr.f_code)

Obviously, as SETUP_WITH explains, the method does get looked up and pushed onto the stack for WITH_CLEANUP to use later. So, even after POP_TOP removes the return value of silly(), its __exit__ is still on the stack.

But there's no way to get at that from Python. Unless you want to start munging the bytecode, or digging apart the stack with ctypes or something, it might as well not exist.

Answer 3

If the context manager is a class and only ever has a single instance, then you could find it on the heap:

import gc

class ConMan(object):
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print "enter %s" % self.name

    def found(self):
        print "You found %s!" % self.name

    def __exit__(self, *args):
        print "exit %s" % self.name


def find_single(typ):
    single = None
    for obj in gc.get_objects():
        if isinstance(obj, typ):
            if single is not None:
                raise ValueError("Found more than one")
            single = obj
    return single



def foo():
    conman = find_single(ConMan)
    conman.found()

with ConMan('the-context-manager'):
    foo()

(Disclaimer: Don't do this)

Answer 4

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.

Answer 5

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!