Create a "With" Block on Several Context Managers

Create a with block on several context managers?

In Python 2.7 and 3.1 and above, you can write:

with A() as X, B() as Y, C() as Z:
    do_something()

This is normally the best method to use, but if you have an unknown-length list of context managers you'll need one of the below methods.

---

In Python 3.3, you can enter an unknown-length list of context managers by using contextlib.ExitStack:

with ExitStack() as stack:
    for mgr in ctx_managers:
        stack.enter_context(mgr)
    # ...

This allows you to create the context managers as you are adding them to the ExitStack, which prevents the possible problem with contextlib.nested (mentioned below).

contextlib2 provides a backport of ExitStack for Python 2.6 and 2.7.

---

In Python 2.6 and below, you can use contextlib.nested:

from contextlib import nested

with nested(A(), B(), C()) as (X, Y, Z):
    do_something()

is equivalent to:

m1, m2, m3 = A(), B(), C()
with m1 as X:
    with m2 as Y:
        with m3 as Z:
            do_something()

Note that this isn't exactly the same as normally using nested with, because A(), B(), and C() will all be called initially, before entering the context managers. This will not work correctly if one of these functions raises an exception.

contextlib.nested is deprecated in newer Python versions in favor of the above methods.

Python nested context manager on multiple lines

Python 3.10 and newer

Starting from Python 3.10, parentheses are allowed, and you can finally do this:

with (
    context1 as a,
    context2 as b
):
    pass

Backslash characters

Two or more physical lines may be joined into logical lines using
backslash characters (\)

(citing the Explicit line joining section)

If you want put context managers on different lines, you can make that work by ending lines with backslashes:

with context1 as a,\
     context2 as b:
    pass

contextlib.ExitStack

contextlib.ExitStack is a

context manager that is designed to make it easy to programmatically
combine other context managers and cleanup functions, especially those
that are optional or otherwise driven by input data.

It's available in Python 3.3 and newer, and allows to enter a variable number of context managers easily. For just two context managers the usage looks like this:

from contextlib import ExitStack

with ExitStack() as es:
    a = es.enter_context(context1)
    b = es.enter_context(context2)

Nesting

It's possible to split a context expression across several nested with statements:

With more than one item, the context managers are processed as if
multiple with statements were nested:

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

suite is equivalent to

with A() as a:
    with B() as b:
        suite

(from The with statement)

Combine two context managers into one

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.

Using different context managers depending on condition

You can store the constructed object in a variable, like:

if some_condition:
    cm = ContextManager(**args)
else:
    cm = OtherContextManager(**other_args)

with cm as contex:
        ... # some block

The above can easily be extended to three possible context managers, etc. You can also decide for example to first "patch" the context manager before "entering" the context.

Although it is common to see a pattern like with foo() as bar:, in fact Python simply evaluates the foo(), obtains that element, and calls .__enter__() on the object. The result of that method is stored in the bar.

So there is nothing "special" about the foo() call, you can use any kind of object on the left side. You can thus for example encapsulate the if-else logic in a separate function, and return the context manager, and then use the variable, or pass context managers as parameters. As long as you use it in a with statement, Python will call .__enter__(..) and .__exit__(..) behind the curtains.

Python multiple context managers in one class

Generally, a very simple way to implement context managers is to use the contextlib module. Writing a context manager becomes as simple as writing a single yield generator. Before the yield replaces the __enter__ method, the object yielded is the return value of __enter__, and the section after the yield is the __exit__ method. Any function on your class can be a context manager, it just needs the be decorated as such. For instance, take this simple ConsoleWriter class:

from contextlib import contextmanager

from sys import stdout
from io import StringIO
from functools import partial

class ConsoleWriter:

    def __init__(self, out=stdout, fmt=None):
        self._out = out
        self._fmt = fmt

    @property
    @contextmanager
    def batch(self):
        original_out = self._out
        self._out = StringIO()
        try:
            yield self
        except Exception as e:
            # There was a problem. Ignore batch commands.
            # (do not swallow the exception though)
            raise
        else:
            # no problem
            original_out.write(self._out.getvalue())
        finally:
            self._out = original_out

    @contextmanager
    def verbose(self, fmt="VERBOSE: {!r}"):
        original_fmt = self._fmt
        self._fmt = fmt
        try:
            yield self
        finally:
            # don't care about errors, just restore end
            self._fmt = original_fmt

    def __getattr__(self, attr):
        """creates function that writes capitalised attribute three times"""
        return partial(self.write, attr.upper()*3)

    def write(self, arg):
        if self._fmt:
            arg = self._fmt.format(arg)
        print(arg, file=self._out)

Example usage:

writer = ConsoleWriter()
with writer.batch:
    print("begin batch")
    writer.a()
    writer.b()
    with writer.verbose():
        writer.c()
    print("before reentrant block")
    with writer.batch:
        writer.d()
    print("after reentrant block")
    print("end batch -- all data is now flushed")

Outputing:

begin batch
before reentrant block
after reentrant block
end batch -- all data is now flushed
AAA
BBB
VERBOSE: 'CCC'
DDD

Can Context Managers run the included block multiple times in Python?

It's not possible. I tried to add multiple yield statements to a context manager, and Python threw a fit. This answer addresses that more, and explains some good alternatives.

This guy examined the bytecode produced, and found that this is not possible. (This guide explains what each bytecode means.)

And this guy shows that the context manager is stored on the heap, which is where classes go, not objects.

Correct way to manage multiple resources with context managers

As mentioned in the comments, ExitStack does exactly this.

A context manager that is designed to make it easy to programmatically combine other context managers

You can simply inherit from ExitStack and call enter_context for each resource you want managed:

class Task(contextlib.ExitStack):
    def __init__(self, res1, res2):
        super().__init__()
        self.a = self.enter_context(Resource(res1))
        self.b = self.enter_context(Resource(res2))

    def run(self):
        print(f'running task with resource {self.a} and {self.b}')

Note there is no need to define your own __enter__ and __exit__ functions, as ExitStack does that for us.

Using it as in the example:

try:
    with Task('foo', 'bar') as t:
        t.run()
except:
    pass

Now when the exception is thrown freeing bar, foo is still freed:

allocating resource foo
allocating resource bar
running task with resource foo and bar
try free resource bar
try free resource foo
freed resource foo

Multiple variables in a 'with' statement?

It is possible in Python 3 since v3.1 and Python 2.7. The new with syntax supports multiple context managers:

with A() as a, B() as b, C() as c:
    doSomething(a,b,c)

Unlike the contextlib.nested, this guarantees that a and b will have their __exit__()'s called even if C() or it's __enter__() method raises an exception.

You can also use earlier variables in later definitions (h/t Ahmad below):

with A() as a, B(a) as b, C(a, b) as c:
    doSomething(a, c)

As of Python 3.10, you can use parentheses:

with (
    A() as a, 
    B(a) as b, 
    C(a, b) as c,
):
    doSomething(a, c)

Split multiple same-level context managers to multiple lines in Python

Line continuations are your friend here...

with \
    open('foo.txt') as foo, \
    open('bar.txt') as bar, \
    open('bla.txt') as bla, \
    open('yada.txt') as yada \
:
    do_something()

This is actually specifically mentioned in PEP-8.

---

Related Topics