CODEWITHSUNDEEP
pythoncontextmanager
Andrew Roberts
Andrew Roberts

Reputation: 800

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

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?

Upvotes: 27

Views: 16766

Answers (4)

A H
A H

Reputation: 2570

In case you're seeing this answer based on the wording of the question, here's a summary of the solution

@contextmanager
def your_custom_context_manager():
    with open(...) as f:
       # do your thing here, e.g. have an `if` statement, or another `with` statement
       yield f  # this is where your new context manager will start from


with your_custom_context_manager() as f:
    # do your stuff

Upvotes: 17

kuzzooroo
kuzzooroo

Reputation: 7408

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!!

Upvotes: 12

Andrew Aylett
Andrew Aylett

Reputation: 40700

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.

Upvotes: 2

jfs
jfs

Reputation: 414179

@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

Upvotes: 24

Related Questions