CodeWalk

@contextmanager装饰器实现上下文管理器

作者：苦行僧 · 2026-05-30 12:55

请解释如何使用@contextmanager装饰器和contextlib.contextmanager将生成器转换为上下文管理器。说明yield前后的代码分别对应__enter__和__exit__的什么部分？如何处理异常并在生成器中重新抛出？给出一个计时器或临时目录切换的实用示例。

回答

苦行僧

@contextmanager将生成器函数变为上下文管理器，yield之前为__enter__逻辑，之后为__exit__逻辑。

from contextlib import contextmanager
import time

@contextmanager
def timer():
    start = time.time()
    yield  # 进入with块的代码在此执行
    elapsed = time.time() - start
    print(f'耗时: {elapsed:.2f}s')

with timer():
    time.sleep(1)

异常处理：若with块内抛出异常，该异常在yield处被抛出生成器。可在yield周围用try/except捕获：

@contextmanager
def safe_file(path):
    f = open(path)
    try:
        yield f
    finally:
        f.close()

注意事项：

  1. 生成器只能yield一次（一次性的）
  2. 若需重复使用，需每次重新调用生成器函数
  3. @asynccontextmanager（Python 3.7+）支持异步上下文管理器

实用示例——临时修改环境变量：

@contextmanager
def temp_env(**kwargs):
    old = {k: os.environ.get(k) for k in kwargs}
    os.environ.update(kwargs)
    try:
        yield
    finally:
        for k, v in old.items():
            if v is None:
                del os.environ[k]
            else:
                os.environ[k] = v