如何实现无需手动写setup的自定义timeit并集成Timer.autorange？

优化方案实现思路 #

1. 无需手动编写setup的实现 #

直接利用timeit.Timer内置的globals参数，传入调用者的全局命名空间即可，完全不需要手动补import或者捕获异常适配。该参数会让测试语句直接从指定的命名空间查找变量、函数，和在当前脚本直接运行的效果完全一致。

2. 集成autorange自动计算循环次数 #

Python 3.6及以上版本的timeit.Timer自带autorange()方法，原生实现了「自动调整循环次数让总运行时长在0.2~2秒区间」的逻辑，直接替换原有的手动循环判断逻辑即可，和IPython的%timeit魔法命令底层逻辑完全对齐。

完整优化后代码 #

import timeit
from typing import Tuple

def timeit_auto(stmt: str = "pass", setup: str = "pass", repeat: int = 3, globals=None) -> Tuple[int, float]:
    """
    模拟%timeit魔法命令行为的自定义计时函数
    :param stmt: 要测试的执行语句
    :param setup: 预执行的配置语句，非必要可留空
    :param repeat: 重复测试的次数，最终取最小值避免系统调度影响
    :param globals: 传入调用方的全局命名空间，填globals()即可自动识别当前作用域的变量/函数
    :return: 单次测试的循环次数、单次执行耗时（单位：微秒）
    """
    t = timeit.Timer(stmt, setup, globals=globals)
    # 自动确定合适的循环次数
    number, _ = t.autorange()
    # 多次重复测试取最小值
    r = t.repeat(repeat, number)
    best = min(r)
    usec = best * 1e6 / number
    return number, usec

使用示例 #

import pandas as pd
import numpy as np

# 构造测试数据
df = pd.DataFrame({
    "col_a": np.random.rand(10000),
    "col_b": np.random.rand(10000)
})
column_name_a = "col_a"

# 直接测试，无需手动写import相关的setup语句
num, timing = timeit_auto(stmt="df[column_name_a]", globals=globals())
print(f"循环次数：{num}，单次耗时：{timing:.2f}微秒")

方案优势 #

• 完全省去手动编写setup导入变量的步骤，只要传入globals=globals()即可自动识别当前作用域的所有变量、函数
• 原生autorange逻辑比手动实现的10次方递增逻辑更合理，适配不同耗时的测试场景
• 不存在原捕获NameError方案的兼容性问题，不会因为异常信息格式变化、变量名包含特殊字符导致导入失败

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