How to write a decorator for generator function in Python?

[justdoit0823]

We can easily write a decorator for simple function, as the following:

import functools
import time

def call_log(func):
    "A simple decorator for recording function call."

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        s_time = time.time()
        ret = func(*args, **kwargs)
        e_time = time.time()
        duration = e_time - s_time
        print('call function %s with duration %f' % (func.__name__, duration))
        return ret

    return wrapper

After, we can decorate other functions with it.

@call_log
def foo():
    print('foo function')

Run the function foo,

In [14]: foo()
foo function
call function %s with duration %f foo 2.09808349609375e-05

But how about the generator function like the following?

@call_log
def bar():
    for i in range(10):
        time.sleep(i)
        yield i

Why the duration is zero?

In [70]: foo()
foo function
call function foo with duration 0.000026

The bar is a generator function, so calling the function will return a generator object, the generator function code actually doesn't execute, but decorator function has finished. So the duration is zero.

Why the result is right?

Because of the decorator function returning generator inside, and tuple executes generator function body and returns right result.

Ah, how can we avoid this situation?

We can yield from the decorated function other than calling it, as the following:

import functools
import time

def call_log(func):
    "A simple decorator for recording function call."

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        s_time = time.time()
        ret = yield from func(*args, **kwargs)
        e_time = time.time()
        duration = e_time - s_time
        print('call function %s with duration %f' % (func.__name__, duration))
        return ret

    return wrapper

With this decorator, the result is correct.

In [73]: tuple(bar())
call function bar with duration 45.020240
Out[73]: (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)

How yield from works?

When yield from is used, it treats the supplied expression as a subiterator. All values produced by that subiterator are passed directly to the caller of the current generator’s methods. Any values passed in with send() and any exceptions passed in with throw() are passed to the underlying iterator if it has the appropriate methods. If this is not the case, then send() will raise AttributeError or TypeError, while throw() will just raise the passed in exception immediately.

When the underlying iterator is complete, the value attribute of the raised StopIteration instance becomes the value of the yield expression. It can be either set explicitly when raising StopIteration, or automatically when the sub-iterator is a generator (by returning a value from the sub-generator).

Changed in version 3.3: Added yield from <expr> to delegate control flow to a subiterator.

The full document is at generator-expressions.

The intresting Python, haha...

[Themis3000]

If you want to control every next() called on the generator instead of passing it through using yield from you can use a decorator class to easily achieve this. Here's an example that allows you to limit the number of iterations in a generator:

class GenLimiter:
  def __init__(self, func):
    self.func = func
  
  def __call__(self, *args, limit=-1, **kwargs):
    self.limit = limit
    self.gen = self.func(*args, **kwargs)
    return self
  
  def __iter__(self):
    return self
  
  def __next__(self):
    if self.limit == 0:
      raise StopIteration()
    self.limit -= 1
    return next(self.gen)


@GenLimiter
def forever():
  num = 0
  while True:
    num += 1
    yield num

# prints numbers 1-5
for item in forever(limit=5):
  print(item)