在Python中，有什么方法可以在调用之前检查函数是否为“生成器函数”？

Question

假设我有两个函数：

def foo():
  return 'foo'

def bar():
  yield 'bar'

第一个是普通函数，第二个是生成器函数。现在我想写一些类似这样的代码：

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

那么，is_generator_function()的简单实现应该是什么样的呢？我可以使用types这个包来测试gen是否是一个生成器，但我希望在调用func()之前就能知道这一点。

现在考虑以下情况：

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

调用goo()会返回一个生成器。我猜想Python的解析器知道goo()函数里面有一个yield语句，我想知道是否有简单的方法可以获取这个信息。

谢谢！

Answer 1

我实现了一个装饰器，它可以处理被装饰函数返回的值或生成的值。基本的思路是：

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

之所以能工作，是因为这个装饰器函数总是会被调用：我们测试的是它的返回值。

---

编辑：根据Robert Lujo的评论，我最后得到了类似这样的代码：

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

Answer 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        
>>> 

你看，主要的区别在于，bar的字节码中至少会包含一个YIELD_VALUE操作码。我建议使用dis模块（当然要把它的输出重定向到一个StringIO实例，然后检查它的getvalue），因为这样可以让你对字节码的变化有一定的抵抗力——操作码的具体数值会变化，但反汇编后的符号值会保持相对稳定;-）。

Answer 3

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

这是关于Python 2.6版本的一些新特性。