如何在调用方忘记对上下文管理器方法使用with块时抛出异常

核心实现思路 #

我们可以自定义带校验逻辑的上下文管理器装饰器，替换原生@contextmanager，通过标记位检测__enter__方法是否被调用，在对象被回收时触发异常提醒或者兜底执行完整逻辑。

完整代码实现 #

import contextlib
import warnings

class CheckedContextManager(contextlib._GeneratorContextManager):
    def __init__(self, gen, func_name):
        super().__init__(gen)
        self._entered = False
        self._func_name = func_name

    def __enter__(self):
        self._entered = True
        return super().__enter__()

    def __del__(self):
        if not self._entered:
            # 方案1：直接抛出异常，强制用户修正使用方式（更推荐）
            raise RuntimeError(f"函数 {self._func_name} 为上下文管理器，必须使用with块调用，正确示例：with {self._func_name}() as val: 你的业务逻辑")
            
            # 方案2：兜底执行，保证after逻辑被触发（适合允许静默修复的场景）
            # warnings.warn(f"检测到{self._func_name}未使用with块调用，已自动执行完整逻辑", UserWarning)
            # try:
            #     self.__enter__()
            # finally:
            #     self.__exit__(None, None, None)

# 自定义装饰器，替换原生@contextmanager
def checked_contextmanager(func):
    def wrapper(*args, **kwargs):
        gen = func(*args, **kwargs)
        return CheckedContextManager(gen, func.__name__)
    return wrapper

使用示例 #

将原来的@contextmanager替换为自定义的@checked_contextmanager即可：

@checked_contextmanager
def foo():
    print("before")
    yield "something"
    print("after")

# 正确用法：正常执行所有逻辑
with foo() as bar:
    print(f"拿到的返回值：{bar}")

# 错误用法：触发异常提醒，或自动执行完整逻辑
foo()

注意事项 #

1. 优先推荐使用抛出异常的方案，从根源上纠正用户的使用习惯，避免业务逻辑未执行到的隐式问题。
2. 兜底执行方案仅适合不需要用户操作yield返回值的场景，如果用户需要在with块内调整返回对象的属性，兜底执行会导致用户的调整逻辑丢失，不符合预期。

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