Python装饰器实战：打造高效性能计时工具

2024-03-31 10:11 由 wang_yb 发表于 #后端开发

在数据分析工作中，针对百万，千万级别的数据进行分析是常有的事情，
因此，分析代码性能的重要性不容忽视，能够有一个方便快速的测试函数性能的方法，
对于我们快速发现性能瓶颈，及时优化，提高项目的开发效率至关重要。

本文介绍如何通过Python装饰器来实现性能计时工具，
帮助我们在不改变现有代码的基础上，随时测试函数的执行时间。

1. 概要

用装饰器来实现这样计时的工具有以下一些好处：
首先，侵入性小，使用装饰器可以非常方便地为函数添加性能计时功能，则无需修改函数的内部代码。
这使得代码更加整洁，也更容易维护。

其次，复用性强，一旦创建了一个性能计时装饰器，就可以将其应用于多个函数，而无需为每个函数单独编写性能计时的代码。
这样不仅提高了代码的效率，也降低了出错的可能性。

最后，是灵活度高，装饰器允许你根据需要定制性能计时的行为，不仅可以打印到终端，也可以根据需求将性能测试结果写入文件或者数据库。

2. 实现计时机制

下面是我目前在用的一个计时装饰器，开发过程中经常用它来看看可能存在性能问题的函数的执行时间。

from functools import wraps
from time import perf_counter


def timeit(loop: int = 1):
    """
    函数执行失败时，重试

    :param loop: 循环执行次数
    :return:
    """

    # 校验参数，参数值不正确时使用默认参数
    if loop < 1:
        loop = 1

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            sum_time: float = 0.0
            for i in range(loop):
                start_time: float = perf_counter()
                ret = func(*args, **kwargs)
                end_time: float = perf_counter()
                sum_time += end_time - start_time

            print(
                f"函数({func.__name__})共执行{loop}次，平均执行时间 {(sum_time/loop):.3f} 秒"
            )
            return ret

        return wrapper

    return decorator

这个装饰器只有一个参数（loop），这个参数可以设置需要反复执行几次待测试的函数（func）。
比如loop设置100，这个装饰器会计算函数（func）执行100次的平均时间。

3. 使用示例

我们用一个模拟的耗时计算函数（compute）来看看这个装饰器的效果。

from decorators import timeit
import time
import random

@timeit(1)
def compute():
    time.sleep(random.random() / 10)
    return 100

if __name__ == "__main__":
    result = compute()
    print(f"{result = }")

从运行结果可以看出，这个timeit装饰器不影响函数的返回值，
不过，这里只执行一次，执行时间存在一定的随机性。

可以将上面代码中的@timeit(1)改成@timeit(100)，再看看执行结果。

执行次数多了之后，平均执行时间开始逼近随机数的中值0.05。
一般性能测试时，都会设置loop这个参数至少大于10，而不会只执行一次。

4. 总结

总之，基于Python装饰器实现的函数性能计时工具具有代码简洁、复用性强、灵活度高、便于性能分析、易于集成等诸多好处。
这些好处使得它成为我们在进行代码性能分析和优化时的有力工具。