I have a decorator function which I want to apply to both normal function and a generator. When applied to the normal function, it works properly. However, when applied to the generator, the iteration loop inside the decorator is executed till the end, but after that the script throws an error:
TypeError: 'NoneType' object is not iterable
and exits the script.
def decor(func):
def wrapper(*args, **kwargs):
func_name = func.__name__
is_generator = "_generator" in func_name
if is_generator:
for item in func(*args, **kwargs):
print(item)
else:
res = func(*args, **kwargs)
print(res)
return wrapper
@decor
def f():
return "a"
@decor
def f_generator():
for i in range(2):
yield "b"
f()
""" Output: a """
for item in f_generator():
print ("Processing item ", item)
"""
Output:
b
b
Traceback (most recent call last):
File "test.py", line 27, in <module>
for item in f_generator():
TypeError: 'NoneType' object is not iterable
"""
Furthermore, when the decorator is applied to the generator, the print ("Processing item ", item) of the external generator call is not executed.
Once I remove the decorator from the generator, I can call the generator and it functions properly.
How do I fix the problem so that I can apply the decorator to the generator and get it working without an error?
Trying to handle the error with an exception takes away the error and the script is executed entirely, but then the print ("Processing item ", item) is still not being executed.
When you add @decorator, f_generator() in for item in f_generator(): is actually decor(f_generator). Since decor() does not yield or return anything, it's not iterable itself, you should add yield item around for item in func(*args, **kwargs):