Python上下文管理器：__exit__异常处理、@contextmanager和ExitStack

with open(...) as f: 这行代码用了十几年，但很多人不知道背后的机制——上下文管理器。更关键的是，自己写一个靠谱的上下文管理器并不简单：__exit__ 里该不该吃掉异常？@contextmanager 的 yield 和 finally 怎么配合？多个资源怎么一起管理？这篇文章把这些问题都讲清楚。

上下文管理器解决什么问题

资源管理的核心要求：用完必须释放，不管有没有异常。不用上下文管理器就得写 try/finally：

python
# 笨办法
f = open('data.txt')
try:
    data = f.read()
finally:
    f.close()

# with 语句：等价但简洁
with open('data.txt') as f:
    data = f.read()
# f.close() 自动调用，即使 read() 抛异常

with 语句保证 __exit__ 一定被调用，省掉 finally 的样板代码。文件、数据库连接、锁、事务，都需要这种保证。

enter 和 exit 协议

上下文管理器是实现 __enter__ 和 __exit__ 的对象：

python
class Timer:
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self  # as 绑定的对象

    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        elapsed = time.perf_counter() - self.start
        print(f"耗时: {elapsed:.4f}s")
        return False  # 不吞异常，继续传播

with Timer() as t:
    time.sleep(1)
# 耗时: 1.0012s

__enter__ 返回值通过 as 绑定。__exit__ 的三个参数是异常信息（没有异常时都是 None）。返回 True 表示吞掉异常，返回 False（或不返回）让异常继续传播。

exit 该不该吞异常

大多数场景不应该吞——返回 False，让调用方处理异常。只有极少数场景需要吞：如 suppress(FileNotFoundError) 这种明确要忽略特定异常的。

python
# 危险：吞掉所有异常，问题难以排查
def __exit__(self, exc_type, exc_val, exc_tb):
    return True  # ❌ 任何异常都被静默忽略

# 正确：只吞特定异常
def __exit__(self, exc_type, exc_val, exc_tb):
    if exc_type is FileNotFoundError:
        return True   # 忽略文件不存在
    return False      # 其他异常继续传播

@contextmanager：更简单的写法

手写 __enter__ 和 __exit__ 容易出错——@contextmanager 装饰器把上下文管理器变成生成器函数，逻辑更清晰：

python
from contextlib import contextmanager

@contextmanager
def timer(name="block"):
    import time
    start = time.perf_counter()
    try:
        yield  # yield 处就是 with 块的代码
    finally:
        elapsed = time.perf_counter() - start
        print(f"{name} 耗时: {elapsed:.4f}s")

with timer("数据处理"):
    process_data()

yield 之前是 __enter__，yield 之后是 __exit__。finally 块保证清理逻辑一定执行。

yield 可以返回值

python
@contextmanager
def temp_directory():
    import tempfile, shutil
    dirpath = tempfile.mkdtemp()
    try:
        yield dirpath  # 返回临时目录路径
    finally:
        shutil.rmtree(dirpath)  # 用完删除

with temp_directory() as tmpdir:
    print(f"临时目录: {tmpdir}")
    # 在 tmpdir 里写文件...
# tmpdir 已被删除

异常处理

@contextmanager 里 yield 抛出的异常会传播到 with 块，但如果在 yield 外面捕获了，就相当于吞掉：

python
@contextmanager
def safe_operation():
    try:
        yield
    except ValueError as e:
        print(f"捕获到 ValueError: {e}")
        # 不 re-raise，异常被吞掉
    finally:
        print("清理完成")

with safe_operation():
    raise ValueError("出错了")
# 输出：捕获到 ValueError: 出错了
#       清理完成
# 程序继续执行，不会崩溃

如果想在 @contextmanager 里记录但不吞掉异常，用 raise 重新抛出，或者不捕获让 finally 执行后自然传播。

contextlib 实用工具

suppress：忽略指定异常

替代 try/except + pass 的惯用法：

python
from contextlib import suppress

# 以前
try:
    os.remove('temp.txt')
except FileNotFoundError:
    pass

# 现在
with suppress(FileNotFoundError):
    os.remove('temp.txt')

可以忽略多种异常：suppress(FileNotFoundError, PermissionError)。

closing：给有 close() 方法的对象加 with 支持

python
from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://example.com')) as response:
    data = response.read()
# response.close() 自动调用

redirect_stdout/redirect_stderr：临时重定向输出

python
from contextlib import redirect_stdout
import io

output = io.StringIO()
with redirect_stdout(output):
    print("这行不会显示在终端")
captured = output.getvalue()  # "这行不会显示在终端
"

适合测试里捕获 print 输出，或者把进度信息写到日志文件而不是终端。

ExitStack：动态管理多个上下文

不确定需要打开多少个资源时用 ExitStack：

python
from contextlib import ExitStack

files = ['a.txt', 'b.txt', 'c.txt']
with ExitStack() as stack:
    handles = [stack.enter_context(open(f)) for f in files]
    # 三个文件都打开了，任何一个打开失败，之前打开的会自动关闭
    for h in handles:
        process(h)
# 三个文件全部自动关闭

也可以用 callback 注册清理函数：

python
with ExitStack() as stack:
    stack.callback(print, "清理完成")
    do_something()
# 无论 do_something 是否抛异常，"清理完成" 都会打印

异步上下文管理器

Python 3.5+ 支持 async with，对应 __aenter__ 和 __aexit__：

python
class AsyncDBConnection:
    async def __aenter__(self):
        self.conn = await create_connection()
        return self.conn

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.conn.close()
        return False

async def main():
    async with AsyncDBConnection() as conn:
        await conn.execute("SELECT 1")

@asynccontextmanager 是异步版本：

python
from contextlib import asynccontextmanager

@asynccontextmanager
async def db_transaction(pool):
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)

async def main():
    async with db_transaction(pool) as conn:
        await conn.execute("INSERT ...")

FastAPI 和 SQLAlchemy 的异步数据库会话就是用这个模式。

实际应用场景

数据库事务

python
@contextmanager
def transaction(conn):
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise

with transaction(conn):
    conn.execute("INSERT ...")
    conn.execute("UPDATE ...")
# 两条语句要么都成功，要么都回滚

临时修改环境变量

python
@contextmanager
def env_var(key, value):
    import os
    old = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = old

with env_var('DATABASE_URL', 'sqlite:///:memory:'):
    run_tests()  # 测试时用内存数据库
# DATABASE_URL 恢复原值

线程锁

threading.Lock 本身就是上下文管理器：

python
import threading

lock = threading.Lock()

with lock:
    # 临界区，自动加锁/释放锁
    update_shared_data()

比 lock.acquire() ... lock.release() 安全——不会因为异常导致锁不释放。