How to decorate iterables with an error handler?

Question

Suppose we have two kinds of methods: one returns a list, the other returns an iterator. So they are very comparable in the sense that both return values are iterable.

I'd like to write a decorator that catches errors inside the iteration. The problem is that the iterator is returned without iteration and so no errors will be caught.

In the below code, the wrapped_properly decorator works around the issue by providing two separate wrappers, a default one ( wrapper ) and one specifically for generator functions ( generatorfunctionwrapper ). The approach feels quite complicated and verbose.

from inspect import isgeneratorfunction
from functools import wraps

def failing_generator():
    for i in range(1, 5):
        if i % 2 == 0:
            print('I dont like even numbers.')
            raise ValueError(i)
        yield i

def wrapped_badly(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except ValueError as err:
            print('Not to worry.')
    return wrapper

def wrapped_properly(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except ValueError as err:
            print('Not to worry.')
    
    @wraps(fn)
    def generatorfunctionwrapper(*args, **kwargs):
        try:
            yield from fn(*args, **kwargs)
        except ValueError as err:
            print('Not to worry.')

    if isgeneratorfunction(fn):
        return generatorfunctionwrapper
    else:
        return wrapper

for x in wrapped_properly(failing_generator)():
    print(x)
# Prints:
# 1
# I dont like even numbers.
# Not to worry.

for x in wrapped_badly(failing_generator)():
    print(x)
# Prints:
# 1
# I dont like even numbers.
# Traceback (most recent call last):
# ...
# ValueError: 2

Is there a better/more pythonic way to do this?

Answer 1

I would suggest returning an iterator no matter what iterable the original function returns.

def wrapped(fn):
    def wrapper(*args, **kwargs):
        try:
            yield from iter(fn(*args, **kwargs))
        except ValueError as err:
            print('Not to worry')

    return wrapper