Nest an iterator of Python context managers in “with”

Question

I have an iterator that returns context managers.

I want a pythonic with statement, that emulates the behaviour of several nested with statements, one for each context manager returned by the iterator.

One could say, I want a generalisation of the (deprecated) contextlib.nested function.

Answer 1

contextlib.nested has two major problems that caused it to be deprecated.

1. first problem is that inner context managers might raise Exceptions during __init__ or __new__, and these exceptions would cause the whole with statement to abort without calling __exit__ of the outer manager.
2. The second problem is more complicated. If one of the inner managers raises an exception and one of the outer managers catches it by returning True in __exit__, the block should still be executed. But in the implementation of nested, it just raises a RuntimeError without executing the block. This problem probably requires a total rewrite of nested.

But it is possible to solve the first problem by just removing one * in the definition of nested! This changes the behaviour such that nested doesn't accept argument lists anymore (which isn't useful anyway because with can handle that already) but only an iterator. I therefore call the new version "iter_nested". The user can then define an iterator that instantiates the context managers during iteration.

An example with a generator:

def contexts():
    yield MyContext1()
    yield MyContext2()

with iter_nested(contexts()) as contexts:
    do_stuff(contexts[0])
    do_other_stuff(contexts[1])

The difference between the codes of the original and my changed version of nested is here:

from contextlib import contextmanager

@contextmanager
--- def nested(*managers):
+++ def iter_nested(mgr_iterator):
    --- #comments & deprecation warning
    exits = []
    vars = []
    --- exc = (None, None, None)
    +++ exc = None # Python 3
    try:
        --- for mgr in managers:
        +++ for mgr in mgr_iterator:
            exit = mgr.__exit__
            enter = mgr.__enter__
            vars.append(enter())
            exits.append(exit)
        yield vars
# All of the following is new and fit for Python 3
except Exception as exception:
    exc = exception
    exc_tuple = (type(exc), exc, exc.__traceback__)
else:
    exc_tuple = (None, None, None)
finally:
    while exits:
        exit = exits.pop()
        try:
            if exit(*exc_tuple):
                exc = None
                exc_tuple = (None, None, None)
        except Exception as exception:
            exception.__context__ = exc
            exc = exception
            exc_tuple = (type(exc), exc, exc.__traceback__)
    if exc:
        raise exc

Answer 2

From the docs:

Developers that need to support nesting of a variable number of context managers can either use the warnings module to suppress the DeprecationWarning raised by [contextlib.nested] or else use this function as a model for an application specific implementation.

The difficult thing about handling multiple context managers is that they interact non-trivially: for example, you might __enter__ the first then raise an exception in __enter__ing the second. These sort of edge cases are precisely what caused nested to be deprecated. If you want to support them, you will have to think very carefully about how you write your code. You may wish to read PEP-0343 for ideas.

Answer 3

This implementation - or something more or less like this, should make what the late contextçlib.nested used to do, but taking care of the already entered contexts if an exception is raised when entering a new context.

Contexts can be passed to it either as a context-protocol object, or as a tuple, where the first member is a called object that will be called with the remainder of the tuple as parameters, in a managed environment:

import sys
import traceback


class NestContext(object):
    def __init__(self, *objects):
        self.objects = objects
    def __enter__(self):
        self.contexts = []
        for obj in self.objects:
            if isinstance(obj, tuple):
                try:
                    obj = obj[0](*obj[1:])
                except Exception, error:
                    self.__exit__(type(error), error, sys.exc_info()[2])
                    raise
            try:
                context = obj.__enter__()
            except Exception, error:
                self.__exit__(type(error), error, sys.exc_info()[2])
                raise   
            self.contexts.append(context)
        return self

    def __iter__(self):
        for context in self.contexts:
            yield context

    def __exit__(self, *args):
        for context in reversed(self.contexts):
            try:
                context.__exit__(*args)
            except Exception, error:
                sys.stderr.write(str(error))

if __name__ == "__main__":
    # example uasage

    class PlainContext(object):
        counter  = 0
        def __enter__(self):
            self.counter = self.__class__.counter
            print self.counter
            self.__class__.counter += 1
            return self
        def __exit__(self, *args):
            print "exiting %d" % self.counter

    with NestContext(*((PlainContext,) for i in range(10))) as all_contexts:
        print tuple(all_contexts)