Лучший способ (если есть) подключить функцию генератора

Question

Я пишу очень простой декоратор, чтобы дать мне базовую отладочную информацию о функции.

from functools import wraps 
from time import perf_counter 

class debug(object): 
    def __init__(self, Time=False, Parameters=False, Doc=False): 
     self.t = Time 
     self.p = Parameters 
     self.d = Doc 

    def __call__(self, func): 
     @wraps(func) 
     def run(*args, **kwargs): 
      params = "" 
      if self.p: 
       params = ", ".join(["{}".format(arg) for arg in args] + ["{}={}".format(k, v) for k, v in kwargs.items()]) 
      print("\n\tDebug output for '{}({})'".format(func.__name__, params)) 
      if self.d: 
       print('\tDocstring: "{}"'.format(func.__doc__)) 
      if self.t: 
       t1 = perf_counter() 
      val = func(*args, **kwargs) 
      if self.t: 
       t2 = perf_counter() 
       print("\tTime Taken: {:.3e} seconds".format(t2 - t1)) 
      print("\tReturn Type: '{}'\n".format(type(val).__name__)) 
      return val 
     return run 

Это все хорошо и хорошо для обычных функций.

@debug(Parameters=True, Time=True, Doc=True) 
def foo(i, j=5): 
    """Raises i to 2j""" 
    for _ in range(j): 
     i **= 2 
    return i 

i = foo(5, j=3) 
# Output: 
""" 
    Debug output for 'foo(5, j=3)' 
    Docstring: "Raises i to 2j" 
    Time Taken: 1.067e-05 seconds 
    Return Type: 'int' 
""" 

Однако генераторы - это совсем другая история.

@debug(Parameters=True, Time=True, Doc=True) 
def bar(i, j=2): 
    """Infinite iterator of increment j""" 
    while True: 
     yield i 
     i += j 

b = bar() # Output occurs here 
next(b) # No output 

Теперь от того, что я закодированы, что вполне ожидаемо, но мне интересно, как я могу подключить метод .__next__() или то, что лучший способ идти об этом есть.

Answer 1

Вы можете просто изменить метод __call__ и возвращают генератор, если генератор приводится в качестве входных данных (добавить import types в верхней части файла):

def __call__(self, f): 
    if isinstance(f, types.GeneratorType): 
     def run_gen(*args, **kwargs): 
      # do pre stuff... 
      for _ in f(*argw, **kwargs): 
       yield _ 
      # do post stuff... 
     return run_gen 
    else: 
     def run(*args, **kwargs): 
      # do pre stuff... 
      r = f(*argw, **kwargs) 
      # do post stuff... 
      return r 
     return run 

Answer 2

Вы не можете заменить function.next, как это значение только для чтения. Но вы можете сделать что-то подобное (см debug_generator функцию):

из functools импорта обертывание импорта инспектировать

class debug(object): 
    def __init__(self, Time=False, Parameters=False, Doc=False): 
     self.t = Time 
     self.p = Parameters 
     self.d = Doc 

    def __call__(self, func): 
     @wraps(func) 
     def debug_generator(func): 
      for i, x in enumerate(list(func)): 
       # here you add your debug statements 
       print "What you want: step %s" % i 
       yield x 
     @wraps(func) 
     def run(*args, **kwargs): 
      params = "" 
      if self.p: 
       params = ", ".join(["{}".format(arg) for arg in args] + ["{}={}".format(k, v) for k, v in kwargs.items()]) 
      print("\n\tDebug output for '{}({})'".format(func.__name__, params)) 
      if self.d: 
       print('\tDocstring: "{}"'.format(func.__doc__)) 
      val = func(*args, **kwargs) 
      print("\tReturn Type: '{}'\n".format(type(val).__name__)) 

      if inspect.isgenerator(val): 
       return debug_generator(val) 
      return val 
     return run 

В основном вы просто получить все значения от генератора вы хотите отлаживать, а затем вы yield их снова, добавив в цикл инструкцию debug.