Generator Context Manager Pitfall

      No Comments on Generator Context Manager Pitfall

Consider the following context manager:

class MyContext(object):
    def __enter__(self):
        print "ENTERED"

    def __exit__(self, exc_type, exc_val, exc_tb):
        print "EXITED", exc_type, exc_val, exc_tb

The following code with print ‘bla’ between ENTERED and EXITED but will also raise the exception:

with MyContext():
    print 'bla'
    raise RuntimeError("ssss")

That’s what I would expect from a context manager – always run the cleanup code.
However, using a context manager built from a generator the behavior is not what I expected:

@contextmanager
def tag(name):
    print '<%s>' % name
    yield
    print '</%s>' % name

The following code will not write the closing tag:

with tag('h1'):
    print 'foo'
    raise RuntimeError('qqq')

So apparently (and of course, it’s documented in the standard), that’s not what happening. But it actually makes sense. If you want to handle the exception that occurs during the contextual execution, you will need to handle an exception as if it’s thrown from the yield statement.

Like so:

@contextmanager
def tag(name):
    print '<%s>' % name
    try:
        yield
    finally:
        print '</%s>' % name

When you implement a context manager using a generator, you do want to know if something happened during the execution and possibly do a different cleanup.
Here’s a good explanation of another scenario where a coroutine’s throw method is used.

Hope you learned something new 🙂

Leave a Reply