Python中如何编写装饰器仅对指定函数进行性能分析

当然可以实现啦！你已经搭好了很棒的代码框架，接下来咱们就用cProfile的内置API把这个装饰器补全，实现只收集指定函数的性能数据～

实现思路 #

咱们不用只局限于命令行调用cProfile，它本身提供了能在代码里直接用的Profile类。核心逻辑就是：创建一个性能分析器实例，用它包裹被装饰函数的执行过程，最后把统计数据统一导出到文件里就行。

完整可运行代码 #

这是补全后的代码，直接就能跑：

import cProfile
import atexit

# 全局性能分析器实例，所有被装饰函数的性能数据会汇总到这里
global_profiler = cProfile.Profile()
# 注册程序退出时自动保存数据，省得咱们手动写保存逻辑
atexit.register(lambda: global_profiler.dump_stats("program.prof"))

def profile_function(func):
    def wrapper(*args, **kwargs):
        # 用全局分析器运行目标函数，自动收集性能数据
        return global_profiler.runcall(func, *args, **kwargs)
    return wrapper


@profile_function
def function1():
    for _ in range(0, 1000):
        _**2


@profile_function
def function2():
    for _ in range(0, 100000):
        _**4


# 不分析这个函数
def function3():
    return


function1()
function2()
function3()

代码说明 #

• 全局的global_profiler会把所有被@profile_function装饰的函数数据汇总到一起，最后导出的program.prof里就只有这两个函数的统计结果，function3不会被记录。
• atexit.register帮咱们在程序正常退出时自动保存数据，不用额外写保存代码。
• global_profiler.runcall(func, *args, **kwargs)这个方法会自动运行函数并收集性能数据，不用手动启用/禁用分析器，省了不少事。

查看结果 #

运行完代码后，还是用你熟悉的命令查看报告就行：

snakeviz program.prof

打开的界面里只会看到function1和function2的性能数据，完全符合你的需求～

可选：每个函数单独保存数据 #

要是你想让每个被装饰的函数生成独立的性能文件，不用汇总到一起，可以把装饰器改成这样：

import cProfile

def profile_function(func):
    def wrapper(*args, **kwargs):
        # 为每个函数创建独立的分析器实例
        local_profiler = cProfile.Profile()
        # 运行函数并收集数据
        result = local_profiler.runcall(func, *args, **kwargs)
        # 以函数名为文件名保存，比如function1.prof、function2.prof
        local_profiler.dump_stats(f"{func.__name__}.prof")
        return result
    return wrapper

运行后会生成两个独立的prof文件，按需查看就行。

小提醒 #

要是你的程序可能异常退出（比如按Ctrl+C中断），atexit的自动保存可能触发不了。这时候可以在主逻辑里加个try/finally块，手动调用global_profiler.dump_stats("program.prof")，确保数据不会丢。

备注：内容来源于stack exchange，提问作者Ohumeronen