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

Answer 2

You write, "I can't do it using @contextmanager - the code after yield call doesn't get executed if you're interrupted by an exception." If you have code that must run you can put it in a try/finally block.

import contextlib

@contextlib.contextmanager
def internal_cm():
    try:
        print "Entering internal_cm"
        yield None
        print "Exiting cleanly from internal_cm"
    finally:
        print "Finally internal_cm"

@contextlib.contextmanager
def external_cm():
    with internal_cm() as c:
        try:
            print "In external_cm_f"
            yield [c]
            print "Exiting cleanly from external_cm_f"
        finally:
            print "Finally external_cm_f"

if "__main__" == __name__:
    with external_cm() as foo1:
        print "Location A"
    print
    with external_cm() as foo2:
        print "Location B"
        raise Exception("Some exception occurs!!")

Output:

Entering internal_cm
In external_cm_f
Location A
Exiting cleanly from external_cm_f
Finally external_cm_f
Exiting cleanly from internal_cm
Finally internal_cm

Entering internal_cm
In external_cm_f
Location B
Finally external_cm_f
Finally internal_cm
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Anaconda\lib\site-packages\spyderlib\widgets\externalshell\sitecustomize.py", line 540, in runfile
    execfile(filename, namespace)
  File "C:\untitled0.py", line 35, in <module>
    raise Exception("Some exception occurs!!")
Exception: Some exception occurs!!

Answer 3

The trouble with what you're doing, is that in using with in your __enter__ call, when you enter your wrapping context manager, you both enter and then leaving the wrapped context managers. If you want write your own context manager that enters the wrapped context managers when you enter the wrapper, then exits them when you leave, you'll have to manually invoke the wrapped context managers' functions. You'll probably also still have to worry about exception safety.