Python中如何识别自定义计时上下文管理器对应的调用代码块

方案1：显式传入代码块名称（最推荐） #

这是最可控、可读性最高的方案，直接给你的Time类构造方法新增一个标识参数，调用时手动传入当前代码块的名称即可：

import time
class Time:
    def __init__(self, block_name: str = "未命名代码块") -> None:
        self.time = 0
        self.block_name = block_name
    def __enter__(self):
        self.start = time.time()
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.time()
        self.time = self.end - self.start
        print(f'代码块 {self.block_name} 执行耗时 {self.time}s')

调用时传入对应标识即可：

with Time("block 1"):
    # block 1
    for i in range(100000):
        for ii in range(100):
            pass
    with Time("block 2"):
        # block 2
        for i in range(1000):
            prime(i)
        with Time("block 3"):
            # block 3
            ''

输出结果会直接对应块名：

代码块 block 3 执行耗时 1.6689300537109375e-06s
代码块 block 2 执行耗时 0.0006444454193115234s
代码块 block 1 执行耗时 0.4997687339782715s

方案2：自动获取调用上下文（无需手动传参） #

如果不想修改现有调用代码，可以借助inspect模块获取调用栈信息，自动拿到代码块所在的文件名、行号、所属函数名来定位：

import time
import inspect
class Time:
    def __init__(self) -> None:
        self.time = 0
        # 获取调用with语句的上层栈帧信息
        call_stack = inspect.stack()[1]
        self.filename = call_stack.filename
        self.lineno = call_stack.lineno
        self.function = call_stack.function
    def __enter__(self):
        self.start = time.time()
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.time()
        self.time = self.end - self.start
        print(f'文件 {self.filename} 第 {self.lineno} 行 函数{self.function} 内的代码块执行耗时 {self.time}s')

这个方案不需要修改调用逻辑，就能直接输出每个计时块的具体位置，适合批量添加计时的场景。

原locals()无效的原因 #

locals()返回的是当前执行作用域的局部变量，你在__exit__方法内调用locals()时，当前作用域是Time类的__exit__方法，自然只能拿到类内部的局部变量，无法获取到外层with块所在作用域的信息。

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