Catch exception thrown in generator caller in Python

Question

I'm trying to catch an exception thrown in the caller of a generator:

class MyException(Exception):
    pass

def gen():
    for i in range(3):
        try:
            yield i
        except MyException:
            print 'handled exception'

for i in gen():
    print i
    raise MyException

This outputs

$ python x.py
0
Traceback (most recent call last):
  File "x.py", line 14, in <module>
    raise MyException
__main__.MyException

when I was intending for it to output

$ python x.py
0
handled exception
1
handled exception
2
handled exception

In retrospect, I think this is because the caller has a different stack from the generator, so the exception isn't bubbled up to the generator. Is that correct? Is there some other way to catch exceptions raised in the caller?

Aside: I can make it work using generator.throw(), but that requires modifying the caller:

def gen():
    for i in range(3):
        try:
            yield i
        except MyException:
            print 'handled exception'
            yield

import sys
g = gen()
for i in g:
    try:
        print i
        raise MyException
    except:
        g.throw(*sys.exc_info())

Answer 1

You may be thinking that when execution hits yield in the generator, the generator executes the body of the for loop, sort of like a Ruby function with yield and a block. That's not how things work in Python.

When execution hits yield, the generator's stack frame is suspended and removed from the stack, and control returns to the code that (implicitly) called the generator's next method. That code then enters the loop body. At the time the exception is raised, the generator's stack frame is not on the stack, and the exception does not go through the generator as it bubbles up.

The generator has no way to respond to this exception.

Answer 2

You might also have gotten confused - as I was just now and arriving here - by the yield in the context managers (contextlib.contextmanager). Their usage can be:

from contextlib import contextmanager


@contextmanager
def mycontext():
    try:
        yield
    except MyException:
        print 'handled exception'

So my solution to a similar case to what you described above is:

def gen():
    for i in range(3):
        yield i


for ii in gen():
    with mycontext():
        print ii
        raise MyException

Which gives the expected output and uses all of the yields.

A bit late to the party here, but maybe someone with similar knots in their minds will find it helpful.

Note: putting the context INSIDE the generator will be the same mistake as using try ... except there!!