Handle an exception thrown in a generator

When a generator throws an exception, it exits. You can't continue consuming the items it generates.

Example:

>>> def f():
...     yield 1
...     raise Exception
...     yield 2
... 
>>> g = f()
>>> next(g)
1
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in f
Exception
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

If you control the generator code, you can handle the exception inside the generator; if not, you should try to avoid an exception occurring.


This is also something that I am not sure if I handle correctly/elegantly.

What I do is to yield an Exception from the generator, and then raise it somewhere else. Like:

class myException(Exception):
    def __init__(self, ...)
    ...

def g():
    ...
    if everything_is_ok:
        yield result
    else:
        yield myException(...)

my_gen = g()
while True:
    try:
        n = next(my_gen)
        if isinstance(n, myException):
            raise n
    except StopIteration:
        break
    except myException as e:
        # Deal with exception, log, print, continue, break etc
    else:
        # Consume n

This way I still carry over the Exception without raising it, which would have caused the generator function to stop. The major drawback is that I need to check the yielded result with isinstance at each iteration. I don't like a generator which can yield results of different types, but use it as a last resort.


I have needed to solve this problem a couple of times and came upon this question after a search for what other people have done.


Throw instead of Raise

One option- which will require refactoring things a little bit- would be to throw the exception in the generator (to another error handling generator) rather than raise it. Here is what that might look like:

def read(handler):
    # the handler argument fixes errors/problems separately
    while something():
        try:
            yield something_else()
        except Exception as e:
            handler.throw(e)
    handler.close()

def err_handler():
    # a generator for processing errors
    while True:
        try:
            yield
        except Exception1:
            handle_exc1()
        except Exception2:
            handle_exc2()
        except Exception3:
            handle_exc3()
        except Exception:
            raise

def process():
    handler = err_handler()
    handler.send(None)  # initialize error handler
    for item in read(handler):
        do stuff

This isn't always going to be the best solution, but it's certainly an option.


Generalized Solution

You could make it all just a bit nicer with a decorator:

class MyError(Exception):
    pass

def handled(handler):
    """
    A decorator that applies error handling to a generator.

    The handler argument received errors to be handled.

    Example usage:

    @handled(err_handler())
    def gen_function():
        yield the_things()
    """
    def handled_inner(gen_f):
        def wrapper(*args, **kwargs):
            g = gen_f(*args, **kwargs)
            while True:
                try:
                    g_next = next(g)
                except StopIteration:
                    break
                if isinstance(g_next, Exception):
                    handler.throw(g_next)
                else:
                    yield g_next
        return wrapper
    handler.send(None)  # initialize handler
    return handled_inner

def my_err_handler():
    while True:
        try:
            yield
        except MyError:
            print("error  handled")
        # all other errors will bubble up here

@handled(my_err_handler())
def read():
    i = 0
    while i<10:
        try:
            yield i
            i += 1
            if i == 3:
                raise MyError()
        except Exception as e:
            # prevent the generator from closing after an Exception
            yield e

def process():
    for item in read():
        print(item)


if __name__=="__main__":
    process()

Output:

0
1
2
error  handled
3
4
5
6
7
8
9

However the downside of this is you have still have to put generic Exception handling inside the generator that might produce errors. It isn't possible to get around this, since raising any exception in a generator will close it.


Kernel of an Idea

It would be nice to have some kind yield raise statement, which allows the generator to continue running if it can after the error was raised. Then you could write code like this:

@handled(my_err_handler())
def read():
    i = 0
    while i<10:
        yield i
        i += 1
        if i == 3:
            yield raise MyError()

...and the handler() decorator could look like this:

def handled(handler):
    def handled_inner(gen_f):
        def wrapper(*args, **kwargs):
            g = gen_f(*args, **kwargs)
            while True:
                try:
                    g_next = next(g)
                except StopIteration:
                    break
                except Exception as e:
                    handler.throw(e)
                else:
                    yield g_next
        return wrapper
    handler.send(None)  # initialize handler
    return handled_inner

After Python 3.3 a code for catching exception from the original generator will be very symple:

from types import GeneratorType


def gen_decorator(func):
    def gen_wrapper(generator):
        try:
            yield from generator  # I mean this line!
        except Exception:
            print('catched in gen_decorator while iterating!'.upper())
            raise

    def wrapper():
        try:
            result = func()

            if isinstance(result, GeneratorType):
                result = gen_wrapper(result)

            return result
        except Exception:
            print('catched in gen_decorator while initialization!'.upper())
            raise

    return wrapper

And example of usage:

@gen_decorator
def gen():
    x = 0
    while True:
        x += 1

        if x == 5:
            raise RuntimeError('error!')

        yield x


if __name__ == '__main__':
    try:
        for i in gen():
            print(i)

            if i >= 10:
                print('lets stop!')
                break
    except Exception:
        print('catched in main!'.upper())
        raise