Wednesday, February 12, 2014

Python: Aggregating Multiple Context Managers

If you make use of context managers you'll eventually run into a situation where you're nesting a number of them in a single with statement.  It can be somewhat unwieldy from a readability point of view to put everything on one line:

with contextmanager1, contextmanager2, contextmanager3, contextmanager4:
    pass


and while you can break it up on multiple lines:

with contextmanager1, \
           contextmanager2, \
           contextmanager3, \
           contextmanager4:
    pass


sometimes that still isn't very readable.  This is more of a problem if you're using the same set of context managers in a number of places.  Ideally you should be able to put the context managers in a variable and use that with however many with statements need them:

handlers = (contextmanager1, contextmanager2, contextmanager3, contextmanager4)
with handlers:
    pass


Of course this doesn't work because handlers is a tuple, not a context manager. This will cause with to throw a exception.  What you can do is create a context manager that aggregates other context managers:

from contextlib import contextmanager
import sys

@contextmanager
def aggregate(handlers):
    for handler in handlers:
        handler.__enter__()
 
    err = None
    exc_info = (None, None, None)
    try:
        yield
    except Exception as err:
        exc_info = sys.exc_info()

    # exc_info get's passed to each subsequent handler.__exit__
    # unless one of them suppresses the exception by returning True
    for handler in reversed(handlers):
        if handler.__exit__(*exc_info):
            err = False
            exc_info = (None, None, None)
    if err:
        raise err

So now you can aggregate all the context managers into one and use that one in the with statement:

handlers = (contextmanager1, contextmanager2, contextmanager3, contextmanager4)
with aggregate(handlers):
    pass


You can build up the list of context managers however you want and use aggregate when using them in a with statement.

§ 

3 comments:

  1. I realized their was a problem with using finally because of the exception suppression so I've updated the context manager to handle that.

    ReplyDelete
  2. Or use the wonderful ExitStack:

    from contextlib import ExitStack

    with ExitStack() as ctx:
    ctx.push(first context manager)
    ctx.push(second context manager)
    ...

    ReplyDelete
    Replies
    1. Cool, I didn't know about ExitStack. I must have completely bypassed it when I read the contextlib docs.

      Delete