怎样利用 sys.settrace 编写一个属于自己的 Python 性能分析工具

在复杂的Python应用中，标准的性能分析工具（如 cProfile）虽然功能强大，但有时我们只需要针对特定的函数或代码块进行轻量级、定制化的时间测量。Python的 sys.settrace 函数提供了一个强大的底层接口，允许我们深入到解释器的执行流程中，从而轻松实现自己的性能分析工具。

什么是 sys.settrace？

sys.settrace(trace_func) 允许您注册一个全局的跟踪函数（trace function）。每当Python解释器遇到特定的事件（如调用函数、执行一行代码、返回或抛出异常）时，这个跟踪函数就会被调用。这是Python调试器和高级分析工具的基础。

跟踪函数的签名必须是 trace_func(frame, event, arg)，其中：

1. frame: 当前正在执行的代码帧对象。
2. event: 发生的事件类型，常用的有 ‘call’ (函数被调用), ‘line’ (执行新的一行), ‘return’ (函数返回), ‘exception’ (抛出异常)。
3. arg: 事件相关的参数（例如，如果是 ‘return’ 事件，arg 就是返回值）。

重要提示： trace_func 必须返回它自己（或者一个新的跟踪函数）以保持跟踪的持续性。

编写自定义的函数级性能分析器

我们将利用 call 事件记录函数的开始时间，并在 return 事件发生时计算并累加函数的执行总时长。

步骤一：创建 Profiler 类

我们使用 time.perf_counter() 来获取高精度的时间，并使用一个字典来存储统计数据。

import sys
import time
from collections import defaultdict

class SimpleFunctionProfiler:
    def __init__(self):
        # 存储统计数据: (filename, func_name) -> {'total_time', 'calls', 'entry_time'}
        self.stats = defaultdict(lambda: {'total_time': 0.0, 'calls': 0, 'entry_time': None})
        self.start_time = time.perf_counter()

    def trace_func(self, frame, event, arg):
        code = frame.f_code
        func_name = code.co_name
        filename = code.co_filename

        # 忽略 profiler 自身代码的跟踪，避免无限递归和性能损耗
        if filename == __file__:
            return self.trace_func

        key = (filename, func_name)
        current_time = time.perf_counter()

        if event == 'call':
            # 记录函数调用开始时间，并增加调用计数
            self.stats[key]['entry_time'] = current_time
            self.stats[key]['calls'] += 1

        elif event == 'return':
            # 计算函数执行时间
            entry_time = self.stats[key]['entry_time']
            if entry_time is not None:
                duration = current_time - entry_time
                self.stats[key]['total_time'] += duration
                # 清除 entry_time，准备迎接下一次调用
                self.stats[key]['entry_time'] = None

        # 必须返回 trace_func 自身，以持续跟踪
        return self.trace_func

    def enable(self):
        """启用跟踪"""
        sys.settrace(self.trace_func)

    def disable(self):
        """禁用跟踪"""
        sys.settrace(None)

    def results(self):
        """打印并格式化结果"""
        print("\n--- Profiling Results ---")
        total_duration = time.perf_counter() - self.start_time
        print(f"Total profiling duration: {total_duration:.4f}s")

        sorted_stats = sorted(
            [(k, v) for k, v in self.stats.items() if v['total_time'] > 0],
            key=lambda item: item[1]['total_time'],
            reverse=True
        )

        print(f"{'Filename:Function':<40} {'Calls':>10} {'Total Time (s)':>20} {'Avg Time (s)':>20}")
        print("-" * 90)

        for (filename, func_name), data in sorted_stats:
            total = data['total_time']
            calls = data['calls']
            avg = total / calls
            # 仅显示文件名，而不是完整路径
            short_filename = filename.split('/')[-1].split('\\')[-1]
            print(f"{short_filename}:{func_name:<38} {calls:>10} {total:>20.6f} {avg:>20.6f}")

步骤二：运行和测试

我们创建一个包含耗时操作的示例代码，并使用我们自定义的 Profiler 进行测量。

# 示例代码
def short_task():
    # 模拟快速操作
    return 1 + 1

def slow_function(n):
    # 模拟耗时操作
    time.sleep(n)
    short_task()

def main_task():
    print("Starting profiling target...")
    slow_function(0.02)
    for i in range(3):
        slow_function(0.005)
    print("Profiling target finished.")

if __name__ == '__main__':
    profiler = SimpleFunctionProfiler()

    # 启用跟踪
    profiler.enable()

    # 执行目标函数
    main_task()

    # 禁用跟踪
    profiler.disable()

    # 显示结果
    profiler.results()

运行结果示例

运行上述代码后，您将看到一个清晰的表格，列出了每个函数被调用的次数、总耗时以及平均耗时（结果会因系统环境而异，但相对时间是准确的）：

Starting profiling target...
Profiling target finished.

--- Profiling Results ---
Total profiling duration: 0.0385s
Filename:Function                             Calls     Total Time (s)         Avg Time (s)
------------------------------------------------------------------------------------------
profiler.py:slow_function                        4             0.035071             0.008768
profiler.py:main_task                            1             0.038541             0.038541
profiler.py:short_task                           4             0.000004             0.000001

总结与注意事项

通过 sys.settrace，我们成功创建了一个轻量级的函数级性能分析工具。然而，使用 sys.settrace 最大的缺点是它会显著增加程序的执行开销（Instrumentation Overhead），因为每次事件发生时都需要调用 Python 函数。对于生产环境中的深度性能分析，可能仍然需要依赖C语言实现的 cProfile。但对于定制化的轻量级检测或教学目的，sys.settrace 是一个极其强大且灵活的工具。