Lets say I have two functions:
def foo():
return \'foo\'
def bar():
yield \'bar\'
The first one is a normal function, and the second i
>>> import inspect
>>>
>>> def foo():
... return 'foo'
...
>>> def bar():
... yield 'bar'
...
>>> print inspect.isgeneratorfunction(foo)
False
>>> print inspect.isgeneratorfunction(bar)
True
I've implemented a decorator that hooks on the decorated function returned/yielded value. Its basic goes:
import types
def output(notifier):
def decorator(f):
def wrapped(*args, **kwargs):
r = f(*args, **kwargs)
if type(r) is types.GeneratorType:
for item in r:
# do something
yield item
else:
# do something
return r
return decorator
It works because the decorator function is unconditionnaly called: it is the return value that is tested.
EDIT: Following the comment by Robert Lujo, I ended up with something like:
def middleman(f):
def return_result(r):
return r
def yield_result(r):
for i in r:
yield i
def decorator(*a, **kwa):
if inspect.isgeneratorfunction(f):
return yield_result(f(*a, **kwa))
else:
return return_result(f(*a, **kwa))
return decorator
>>> def foo():
... return 'foo'
...
>>> def bar():
... yield 'bar'
...
>>> import dis
>>> dis.dis(foo)
2 0 LOAD_CONST 1 ('foo')
3 RETURN_VALUE
>>> dis.dis(bar)
2 0 LOAD_CONST 1 ('bar')
3 YIELD_VALUE
4 POP_TOP
5 LOAD_CONST 0 (None)
8 RETURN_VALUE
>>>
As you see, the key difference is that the bytecode for bar
will contain at least one YIELD_VALUE
opcode. I recommend using the dis
module (redirecting its output to a StringIO instance and checking its getvalue
, of course) because this provides you a measure of robustness over bytecode changes -- the exact numeric values of the opcodes will change, but the disassembled symbolic value will stay pretty stable;-).