Writing a context manager in Python that itself uses a with statement

Question

I'm trying to write a context manager that uses other context managers, so clients don't need to know the whole recipe, just the interface I'm presenting. I can't do it using @contextmanager - the code after yield call doesn't get executed if you're interrupted by an exception, so I need to use a class-based manager.

Here's a little example script:

from contextlib import contextmanager
import pprint

d = {}

@contextmanager
def simple(arg, val):
    print "enter", arg
    d[arg] = val
    yield
    print "exit", arg
    del d[arg]

class compl(object):
    def __init__(self, arg, val):
        self.arg=arg
        self.val=val

    def __enter__(self):
        with simple("one",1):
            with simple("two",2):
                print "enter complex", self.arg
                d[self.arg] = self.val

    def __exit__(self,*args):
        print "exit complex", self.arg
        del d[self.arg]

print "before"
print d
print ""

with compl("three",3):
    print d
    print ""

print "after"
print d
print ""

That outputs this:

before
{}

enter one
enter two
enter complex three
exit two
exit one
{'three': 3}

exit complex three
after
{}

I want it to output this:

before
{}

enter one
enter two
enter complex three
{'one': 1, 'three': 3, 'two': 2}

exit complex three
exit two
exit one
after
{}

Is there any way to tell a class-based context manager to wrap itself with other context managers?

Answer 1

@contextmanager
def compl(arg, val):
    with simple("one",1):
        with simple("two",2):
            print "enter complex", arg 
            try:
                d[arg] = val
                yield
            finally:
                del d[arg]
                print "exit complex", arg