what is the pythonic way to inherit context manager

Question

Python teaches us to do cleanup on objects using __enter__ and __exit__. What if i need to create an object which uses objects must use context managers? Imagine this:

from database1 import DB1
from database2 import DB2

Normally, they would be used as such:

with DB1() as db1, DB2() as db2:
    db1.do_stuff()
    db2.do_other_stuff()

Whatever happens, db1 and db2 will both run their __exit__ function, and clean up the connection, flush, etc.

When i'd put all of this in a class, how would i do it? Is this right? This is obviously not right, the context manager for db1 and db2 runs at the end of the block, as pointed in comments.

class MyApp(object):
    def __enter__(self):
        with DB1() as self.db1, DB2() as self.db2:
            return self
    def __exit__(self, type, value, traceback):
        self.db1.__exit__(self, type, value, traceback)
        self.db2.__exit__(self, type, value, traceback)

I even considered doing something like this: This looks like a good idea, actually (after some cleanup):

class MyApp(object):
    def __init__(self):
        self.db1 = DB1()
        self.db2 = DB2()
    def __enter__(self):
        self.db1.__enter__()
        self.db2.__enter__()
        return self
    def __exit__(self, type, value, traceback):
        try:
            self.db1.__exit__(self, type, value, traceback)
        except:
            pass
        try:
            self.db2.__exit__(self, type, value, traceback)
        except:
            pass

EDIT: Fixed the code.

Answer 1

I would go with the second solution but also handle database errors:

import sys

class MyApp(object):
    def __init__(self):
        self.db1 = DB1()
        self.db2 = DB2()
    def __enter__(self):
        self.db1.__enter__()
        try:
            self.db2.__enter__()
        except:
            self.db1.__exit__(None, None, None) # I am not sure with None
            raise
        return self
    def __exit__(self, type, value, traceback):
        try:
            self.db1.__exit__(self, type, value, traceback)
        finally:
            self.db2.__exit__(self, type, value, traceback)

The first one calls __exit__ in __enter__ because of with - so, does not work.

EDIT: Also check out the answer by @Ming. In many cases it is shorter.

Answer 2

Most context managers can just be written with using the @contextmanager decorator. You write a function with one yield, before the yield is your 'enter' function and after the yield is your 'exit' function. Due to the way generators are implemented if a yield is in a with statement then the with statement does not exit at the yield.

eg.

from contextlib import contextmanager

class SomeContextManager:
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        print("enter", self.name)
        return self
    def __exit__(self, ex_type, value, traceback):
        print("exit", self.name)

class SomeContextManagerWrapper:
    def __init__(self, *context_managers):
        self.context_managers = context_managers
    @property
    def names(self):
        return [cm.name for cm in self.context_managers]

@contextmanager
def context_manager_combiner():
    print("context_manager_combiner entering")
    with SomeContextManager("first") as a, SomeContextManager("second") as b:
        yield SomeContextManagerWrapper(a, b)
    print("context_manager_combiner exiting")

with context_manager_combiner() as wrapper:
    print("in with statement with:", wrapper.names)

outputs:

context_manager_combiner entering
enter first
enter second
in with statement with: ['first', 'second']
exit second
exit first
context_manager_combiner exiting

Answer 3

Depends on what you are trying to achieve overall. One possibility is to construct the individual context managers, then combine them with the standard library's contextlib.nested. This will give you a single object which behaves like your example MyApp but utilizes the existing standard library in a DRY (Don't Repeat Yourself) fashion.