如何计算Python通用生成器函数每次yield的执行耗时？

我尝试计算任意生成器函数中每次yield的执行耗时，但始终找不到合适方法。目前最接近的方案是编写如下包装器：

def catchtime_wrapper_for_method(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return_val = func(*args, **kwargs)
        if isinstance(return_val, types.GeneratorType):
            for g in return_val:
                start_time = perf_counter()
                yield g
                end_time = perf_counter()
                print(f"Iteration time: {end_time - start_time:.6f} seconds")
        else:
            return return_val
    return wrapper

调用方式如下：

for i in catchtime_wrapper_for_method(my_gen_function)(args...):
   # 此处处理i

my_gen_function会yield IO操作及相关计算的结果，其中集成的tqdm已验证准确，显示每次迭代耗时约5秒。但上述包装器输出的每次迭代耗时仅约1秒，推测该时间未包含my_gen_function的执行耗时，实际准确耗时应为约6秒。我需要通过包装器输出准确耗时，同时保持my_gen_function的生成器行为不变。

我还尝试过以下方法（可能存在操作失误）：

• yield from
• 装饰器
• 利用__enter__和__exit__的catchtime类
• 使用next()方法

请问这一偏差产生的原因是什么？

偏差产生的原因 #

你的包装器计时逻辑完全搞反了顺序：

• 当用for g in return_val遍历生成器时，Python会隐式调用next(gen)，这会触发生成器运行到下一个yield语句——也就是my_gen_function里完成所有IO、计算逻辑，直到产出值的全过程，这部分才是耗时的核心。
• 而你的代码是在生成器已经产出值之后才开始计时，yield g只是把已经生成好的值传递给调用方，这个过程几乎不耗时，所以统计出的1秒只是值传递的时间，完全没包含my_gen_function里生成该值的5秒耗时。

修正方案 #

要准确统计每次迭代的完整耗时，必须在触发生成器执行之前开始计时，也就是在调用next()之前启动计时器。修改后的包装器如下：

import types
from functools import wraps
from time import perf_counter

def catchtime_wrapper_for_method(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return_val = func(*args, **kwargs)
        if isinstance(return_val, types.GeneratorType):
            gen = return_val
            while True:
                try:
                    # 触发生成器执行前开始计时
                    start_time = perf_counter()
                    g = next(gen)  # 这里会让生成器运行到下一个yield，包含所有前置IO/计算
                    end_time = perf_counter()
                    print(f"Iteration time: {end_time - start_time:.6f} seconds")
                    yield g
                except StopIteration:
                    # 生成器耗尽时退出循环
                    break
        else:
            return return_val
    return wrapper

这个版本的逻辑是：

1. 手动控制生成器的迭代过程，每次调用next(gen)前启动计时
2. next(gen)会驱动my_gen_function执行到下一个yield，完成所有耗时操作并产出值
3. 获取到值后立即停止计时，此时统计的时间就是从生成器开始执行到产出值的完整耗时
4. 最后把产出的值yield给调用方，保持原生成器的行为不变

内容的提问来源于stack exchange，提问作者Benny Semyonov