How to remove a context manager from an ExitStack

Question

I have a long-running process called an Updater, that has updates (to an ETL system) submitted to it. The updates have resource requirements that are managed by adding a context manager to the Updater's ExitStack. Some updates will include new configuration, which means that effected resources must be released from the stack and a newly configured version of the resource will be added. I need something like:

with ExitStack() as stack:
    ctx_manager = open("file.txt")
    f = stack.enter_context(ctx_manager)
    ...
    ctx_pop(ctx_manager, stack)  # remove the given context manager from the stack

Below is an example of something I've gotten to work, but it relies on accessing protected members. I was hoping there might be a less 'dirty' solution than this:

def ctx_pop(cm, stack):
    for item in stack._exit_callbacks:
        if item.__self__ is cm:
            break
    else:
        raise KeyError(repr(cm))
    stack._exit_callbacks.remove(item)
    item(None, None, None)

Edit: Added known solution

Protocol

Looking at the source code ExitStack uses a deque to store wrappers for the contexts .__exit__ method so you would need to be able to identify the wrapper by the original context manager which, as far as I know, isn't possible. You may need to reinvent (at least partially) the functionality of ExitStack to be able to remove contexts ahead of time. – Archeozoic

@TadhgMcDonald-Jensen thank you. I discovered the same thing inspecting ExitStack objects in a Jupyter Notebook. They have a deque of closures who's __self__ attribute is a context manager. I'll put some code up as an known solution, but I was hoping for a less 'hack' solution. I wanted to see if there was a cleaner solution before I submit to Python Ideas or add it myself (which I've never done before). – Protocol

Answer 1

You have to extend ExitStack with your own pop-method:

from contextlib import ExitStack
from collections import deque

class ExitStackWithPop(ExitStack):
    def pop(self, cm):
        callbacks = self._exit_callbacks
        self._exit_callbacks = deque()
        found = None
        while callbacks:
            cb = callbacks.popleft()
            if cb.__self__ == cm:
                found = cb
            else:
                self._exit_callbacks.append(cb)
        if not found:
            raise KeyError("context manager not found")
        found(None, None, None)

Barye

it is not _ExitStack__self__ it is just __self__, attributes with two trailing underscores are not private (just think of all the magic method names) Also you don't need to recreate the deque, you can just iterate through the items and .remove the one that has .__self__ == cm – Archeozoic

Thanks! That's very close to the solution I'm using right now, but I like the idea of subclassing. It seem cleaner that way. – Protocol

Bleh. You're reaching into the implementation details here. – Hobby

Since this changes at contextlib it must be if cb[1].__self__ == cm: to work. (Anyway its a very low level dependency to contextlib) – Hilliard

it is odd that this functionality is not included in ExitStack. It really should. Clearly there are use-cases where one needs to manage objects dynamically. Have you considered submitting an enhacement request? – Uniaxial

I just submitted an enhancement request bugs.python.org/issue45184 – Uniaxial

Answer 2

contextlib.ExitStack only supports exiting all context managers at once. You can't pop context managers individually. If you want to do this, you should keep track of your context managers with something other than an ExitStack.

Hobby

Reimplementing the functionality similar to contextlib.ExitStack (keeping track of context managers) is certainly not easy. – Vin

Answer 3

A simple option that doesn't rely on poking around with the internals of ExitStack:

from contextlib import ExitStack, AbstractContextManager
from dataclasses import dataclass, field

@dataclass
class CtxStack:
    ctx_stack: list[ExitStack] = field(default_factory=list, init=False)

    def enter[ctxT](self, ctx: AbstractContextManager[ctxT]) -> ctxT:
        mgr = ExitStack()
        result = mgr.enter_context(ctx)
        self.ctx_stack.append(mgr)
        return result
    
    def pop(self) -> None:
        self.ctx_stack.pop().pop_all()

Zelda