意外 decorator/generator 行为

Question

我创建了一个计时器装饰器，但如果装饰函数是一个生成器，它就不起作用。

import numpy as np
from time import time
from collections import Counter


def timer(f):
    def inner(*args, **kwargs):
        start = time()
        try:
            res = f(*args, **kwargs)
        except Exception as e:
            end = time()
            timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
            raise e
        end = time()
        timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
        return res
    return inner
timer.counter = Counter()


class AA:
    @timer
    def __init__(self):
        a = np.array(range(1_000_000))

    @timer
    def __iter__(self):
        a = np.array(range(1_000_000))
        yield 'a'

    @timer
    def normal_fun(self):
        a = np.array(range(1_000_000))

    @timer
    def fun_with_yield(self):
        a = np.array(range(1_000_000))
        yield 'a'


a = AA()
for i in a:
    pass
a.normal_fun()
a.fun_with_yield()
print(timer.counter)

输出：

Counter({'main.init': 0.10380005836486816, 'main.normal_fun': 0.10372400283813477, 'main.iter': 0.0, 'main.fun_with_yield': 0.0})

为什么生成器函数的时间等于 0.0，我该如何解决？

Answer 1

最后我通过为生成器创建另一个装饰器达到了预期的效果。装饰器本身就是一个生成器，计算每次迭代所花费的时间。

def generator_timer(f):
    def inner(*args, **kwargs):
        start = time()
        try:
            res = f(*args, **kwargs)
            for i in res:
                end = time()
                timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
                yield i
                start = time()
        except Exception as e:
            end = time()
            timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
            raise e
        end = time()
        timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
    return inner


class AA:
    @timer
    def __init__(self):
        a = np.array(range(10_000_000))

    @generator_timer
    def __iter__(self):
        a = np.array(range(10_000_000))
        yield 'a'
        yield 'b'

    @timer
    def normal_fun(self):
        a = np.array(range(10_000_000))

    @generator_timer
    def fun_with_yield(self):
        a = np.array(range(10_000_000))
        yield a

Counter({'main.init': 1.0399727821350098, 'main.iter': 1.0183088779449463, 'main.fun_with_yield': 1.0168907642364502, 'main.normal_fun': 1.0156745910644531})

意外 decorator/generator 行为

Unexpected decorator/generator behavior

python

generator

python-decorators