如何利用上下文管理器实现简洁高效的代码执行计时方案？

需求说明 #

• 可快速添加至代码，不改变原有结构与逻辑；
• 支持对单行代码或代码块进行计时；
• 实现方式足够优雅简洁，可长期保留在代码中。

尝试过timeit模块与time模块后均未满足需求，最终通过上下文管理器的__enter__与__exit__魔术方法实现该方案。核心思路为：在上下文管理器启动（__enter__）时记录起始时间，待上下文块中的代码执行完成后，在__exit__方法中记录结束时间，计算并处理计时结果（打印或返回）。

完整实现代码 #

import time


class ExecutionTimer:

    def __init__(self, message='Execution time', unit='s'):
        self.message = message
        self.unit = unit
        self.start_time = None
        self.end_time = None
        self.result = None

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.result = self.end_time - self.start_time
        self.print_time()

    def print_time(self):
        elapsed_time = self.result
        if self.unit == 'ms':
            elapsed_time *= 1000  # convert to milliseconds
        elif self.unit == 'us':
            elapsed_time *= 1e6  # convert to microseconds
        elif self.unit == 'ns':
            elapsed_time *= 1e9  # convert to nanoseconds

        print(f"{self.message}: {elapsed_time}{self.unit}")


if __name__ == '__main__':
    
    start_time = time.time()
    time.sleep(1)
    end_time = time.time()
    print(f"Execution (s): {end_time - start_time}s")

    with ExecutionTimer(message='Execution (s)', unit='s'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (ms)", unit='ms'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (us)", unit='us'):
        time.sleep(1)

    with ExecutionTimer(message="Execution (ns)", unit='ns'):
        time.sleep(1)

    # or just capture the value
    with ExecutionTimer(message="Execution (ns)", unit='ns') as my_timer:
        time.sleep(1)

    # notice we are not in the context manager any more
    print(my_timer.result)
    print(my_timer.result)

执行输出 #

Execution (s): 1.0000789165496826s
Execution (s): 1.0000693798065186s
Execution (ms): 1000.067949295044ms
Execution (us): 1000072.0024108887us
Execution (ns): 1000071287.1551514ns
Execution (ns): 1000077962.8753662ns
1.0000779628753662
1.0000779628753662

Process finished with exit code 0

优化方向 #

• 集成日志系统，按指定级别输出计时信息；
• 添加dry_run参数，控制计时结果的打印逻辑；
• 支持自定义结果处理逻辑（如写入文件、统计分析等）。

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