实现带参数的时间统计装饰器

题目

实现带参数的时间统计装饰器

信息

• 类型：问答
• 难度：⭐⭐

考点

装饰器定义与使用,带参数装饰器,函数包装,时间处理

快速回答

实现一个带参数的装饰器 @timeit(unit='ms') 用于统计函数执行时间，支持自定义时间单位（秒/毫秒/微秒）。核心要点：

• 使用三层嵌套函数结构实现带参装饰器
• 利用 time.perf_counter() 高精度计时
• 通过 @functools.wraps 保留原函数元信息
• 实现时间单位自动转换（s/ms/us）
• 处理无效单位异常

原理说明

带参数装饰器本质是三层嵌套函数：
1. 外层接收装饰器参数
2. 中层接收被装饰函数
3. 内层实现具体包装逻辑
通过 time.perf_counter() 获取高精度时间戳，差值即为执行耗时。

代码示例

import time
import functools

def timeit(unit='ms'):
    # 验证时间单位有效性
    valid_units = {'s', 'ms', 'us'}
    if unit not in valid_units:
        raise ValueError(f"Invalid unit '{unit}'. Use 's', 'ms' or 'us'")

    def decorator(func):
        @functools.wraps(func)  # 保留元信息
        def wrapper(*args, **kwargs):
            # 计时开始
            start = time.perf_counter()
            result = func(*args, **kwargs)  # 执行原函数
            # 计算耗时并转换单位
            duration = time.perf_counter() - start

            # 单位转换
            if unit == 'ms':
                duration *= 1000
            elif unit == 'us':
                duration *= 1_000_000

            print(f"{func.__name__} executed in {duration:.4f} {unit}")
            return result
        return wrapper
    return decorator

# 使用示例
@timeit(unit='ms')
def calculate_sum(n):
    return sum(range(n))

calculate_sum(1000000)  # 输出: calculate_sum executed in 25.4321 ms

最佳实践

• 保留元信息：始终使用 @functools.wraps 避免调试问题
• 高精度计时：perf_counter() 比 time.time() 更适合性能测量
• 参数校验：装饰器参数需做有效性检查
• 线程安全：嵌套函数确保装饰器状态隔离

常见错误

• 缺少 functools.wraps：导致被装饰函数 __name__ 等属性丢失
• 单位转换错误：忘记时间单位是秒的倍数关系（1s=1000ms）
• 装饰器嵌套错误：三层函数缺少任意一层都会导致调用异常
• 计时位置错误：start 必须在调用原函数前获取

扩展知识

• 类实现装饰器：通过实现 __call__ 和 __init__ 方法
• 多装饰器顺序：@A @B def foo() 等价于 foo = A(B(foo))
• 异步函数支持：使用 async def 并配合 asyncio.coroutine
• 性能优化：生产环境建议结合 logging 模块替代 print
• 动态启停：可通过附加参数控制是否启用计时功能