python - Combine two context managers into one

Question

I use Python 2.7 and I know that I can write this:

with A() as a, B() as b:
    do_something()

I want to provide a convenience helper which does both. The usage of this helper should look like this:

with AB() as ab:
    do_something()

Now AB() should do both: Create context A() and create context B().

I have no clue how to write this convenience helper

Answer 1

Don't re-invent the wheel; this is not as simple as it looks.

Context managers are treated as a stack, and should be exited in reverse order in which they are entered, for example. If an exception occurred, this order matters, as any context manager could suppress the exception, at which point the remaining managers will not even get notified of this. The __exit__ method is also permitted to raise a different exception, and other context managers then should be able to handle that new exception. Next, successfully creating A() means it should be notified if B() failed with an exception.

Now, if all you want to do is create a fixed number of context managers you know up front, just use the @contextlib.contextmanager decorator on a generator function:

from contextlib import contextmanager

@contextmanager
def ab_context():
    with A() as a, B() as b:
        yield (a, b)

then use that as:

with ab_context() as ab:

If you need to handle a variable number of context managers, then don't build your own implementation; use the standard library contextlib.ExitStack() implementation instead:

from contextlib import ExitStack

with ExitStack() as stack:
    cms = [stack.enter_context(cls()) for cls in (A, B)]

    # ...

The ExitStack then takes care of correct nesting of the context managers, handling exiting correctly, in order, and with the correct passing of exceptions (including not passing the exception on when suppressed, and passing on new-ly raised exceptions).

If you feel the two lines (with, and separate calls to enter_context()) are too tedious, you can use a separate @contextmanager-decorated generator function:

from contextlib import ExitStack, contextmanager

@contextmanager
def multi_context(*cms):
    with ExitStack() as stack:
        yield [stack.enter_context(cls()) for cls in cms]

then use ab_context like this:

with multi_context(A, B) as ab:
    # ...

For Python 2, install the contextlib2 package, and use the following imports:

try:
    from contextlib import ExitStack, contextmanager
except ImportError:
    # Python 2
    from contextlib2 import ExitStack, contextmanager

This lets you avoid reinventing this wheel on Python 2 too.

Whatever you do, do not use contextlib.nested(); this was removed from the library in Python 3 for very good reasons; it too did not implement handling entering and exiting of nested contexts correctly.