В python есть ли способ проверить, является ли функция "функцией генератора" перед ее вызовом?

Предположим, что у меня есть две функции:

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, и мне интересно, легко ли получить эту информацию.

Спасибо!

Ответы

Ответ 1

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

Новый в версии Python 2.6

Ответ 2

На самом деле, мне интересно, насколько полезен такой гипотетический is_generator_function(). Рассмотрим:

def foo():
    return 'foo'
def bar():
    yield 'bar'
def baz():
    return bar()
def quux(b):
    if b:
        return foo()
    else:
        return bar()

Что должно is_generator_function() возвращаться для baz и quux? baz() возвращает генератор, но не сам по себе, а quux() может возвращать генератор или не может.

Ответ 3

>>> 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), потому что это обеспечивает вам определенную степень надежности при изменении байт-кода - точные числовые значения кодов операций будут меняться, но разобранное символическое значение останется довольно стабильным; -).

Ответ 4

Я реализовал декоратор, который перехватывает украшенную функцию return/yielded value. Его основные положения:

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

Это работает, потому что функция декоратора называется unconditionnaly: это проверенное возвращаемое значение.

РЕДАКТИРОВАТЬ: Следуя комментарию Роберта Луджо, у меня получилось что-то вроде:

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