In python is there a way to check if a function is a "generator function" before calling it?

Lets say I have two functions:

def foo():
  return 'foo'

def bar():
  yield 'bar'

The first one is a normal function, and the second is a generator function. Now I want to write something like this:

def run(func):
  if is_generator_function(func):
     gen = func()
     gen.next()
     #... run the generator ...
  else:
     func()

What will a straightforward implementation of is_generator_function() look like? Using the types package I can test if gen is a generator, but I wish to do so before invoking func().

Now consider the following case:

def goo():
  if False:
     yield
  else:
     return

An invocation of goo() will return a generator. I presume that the python parser knows that the goo() function has a yield statement, and I wonder if it possible to get that information easily.

Thanks!

Solution 1:

>>> import inspect
>>> 
>>> def foo():
...   return 'foo'
... 
>>> def bar():
...   yield 'bar'
... 
>>> print inspect.isgeneratorfunction(foo)
False
>>> print inspect.isgeneratorfunction(bar)
True

• New in Python version 2.6

Solution 2:

>>> 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;-).