上周帮一个朋友排查一个生产环境的数据库连接泄漏问题，查了半天发现罪魁祸首是一个自定义的上下文管理器——__exit__ 方法里少写了一行 close()。服务器跑了三天，连接池被榨干了。

这件事让我意识到，很多人对 Python 的 with 语句停留在"打开文件记得用 with"这个层面，但真到自己写上下文管理器的时候，各种细节上的坑就冒出来了。

这篇文章总结了我这几年在上下文管理器上遇到的几个典型案例，以及一些不那么常见但很实用的用法。

一：__exit__ 吞掉了异常

最常见的自定义上下文管理器写法是这样：

class DatabaseConnection:
    def __init__(self, conn_string):
        self.conn_string = conn_string
        self.conn = None

    def __enter__(self):
        self.conn = create_connection(self.conn_string)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.close()

看起来没什么问题对吧？但如果 create_connection 抛了异常呢？这时候 self.conn 还是 None，__exit__ 里调用 self.conn.close() 就会抛出 AttributeError。

这还不是最糟的。更隐蔽的问题是：如果 __exit__ 返回了 True，Python 会认为异常已经被处理了，不会往外抛。比如：

class DangerousContext:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 做了一个操作，顺便返回了 True
        cleanup()
        return True  # 所有异常都被静默吞掉了！

__exit__ 返回 True 意味着"这个异常我处理好了，你不用管了"。如果你在 __exit__ 里不小心返回了一个 true 值，那代码块里抛出的任何异常都会凭空消失，调试起来非常痛苦。

正确做法：__exit__ 方法里要么不返回（默认 None），要么只在确实处理了异常时才返回 True。清理逻辑用 try/finally 包裹：

def __exit__(self, exc_type, exc_val, exc_tb):
    if self.conn:
        try:
            self.conn.close()
        except Exception:
            pass  # 至少记录一下日志
    # 不返回 True，让异常正常传播

二：contextlib.contextmanager 里 yield 之后的代码不一定会执行

用 @contextmanager 装饰器写上下文管理器很方便，但很多人没搞清楚它的执行模型：

from contextlib import contextmanager

@contextmanager
def managed_resource():
    resource = acquire_resource()
    try:
        yield resource
    finally:
        resource.release()

看起来总算安全了吧？finally 块总会执行的。问题是，如果你忘了写 try/finally：

@contextmanager
def broken_resource():
    resource = acquire_resource()
    yield resource
    resource.release()  # 如果 with 块里抛了异常，这行不会执行！

如果 with 代码块里抛出了异常，yield 会把异常往外抛，yield 后面的代码根本不会执行。资源就泄露了。

更隐蔽的是，如果你在 yield 后面写了 return：

@contextmanager
def sneaky_bug():
    resource = acquire()
    yield resource
    return "cleanup done"  # 这个 return 值会去哪里？

这个返回值会被生成器的 StopIteration 捕获，然后被 @contextmanager 忽略。不会报错，但也不会有任何效果。你的"清理完毕"返回值就悄无声息地蒸发了。

记住一条铁律：用 @contextmanager 时，yield 必须放在 try 块里，清理代码必须放在 finally 里。没有例外。

三：嵌套上下文管理器的缩进地狱

看看这种代码你写过没有：

with open('config.yaml') as f:
    config = yaml.safe_load(f)
    with create_db_connection(config['db']) as conn:
        with conn.cursor() as cursor:
            cursor.execute("SELECT * FROM users")
            results = cursor.fetchall()
            with open('output.csv', 'w') as out:
                writer = csv.writer(out)
                writer.writerows(results)

四层缩进，看着就头疼。Python 3.10 之后可以用括号把多个上下文管理器写在一行：

with (
    open('config.yaml') as f,
    create_db_connection(yaml.safe_load(f)['db']) as conn,
):
    # 等一下，yaml.safe_load(f) 在第二个 with 表达式里？
    # 这时候 f 还没绑定呢...

实际上这里有个先后顺序的问题——所有表达式在进入 with 之前按顺序求值，但 as 绑定的变量要等对应的 __enter__ 返回后才能用。所以第二个表达式里访问 f 是拿不到的。

更稳妥的做法是用 contextlib.ExitStack：

from contextlib import ExitStack

with ExitStack() as stack:
    f = stack.enter_context(open('config.yaml'))
    config = yaml.safe_load(f)
    conn = stack.enter_context(create_db_connection(config['db']))
    cursor = stack.enter_context(conn.cursor())
    out = stack.enter_context(open('output.csv', 'w'))

    cursor.execute("SELECT * FROM users")
    writer = csv.writer(out)
    writer.writerows(cursor.fetchall())

ExitStack 会按后进先出的顺序自动清理所有注册的上下文。而且它还有一个很实用的能力——你可以动态地往里面加上下文：

with ExitStack() as stack:
    conn = stack.enter_context(get_db())

    for table in tables_to_backup:
        # 根据运行时条件决定打开哪些文件
        f = stack.enter_context(open(f'backup_{table}.csv', 'w'))
        dump_table(conn, table, f)

这在处理数量不定的文件或连接时非常方便。

四：异步上下文管理器忘记用 async with

现在用 asyncio 的项目越来越多，很多库提供了异步版本的上下文管理器。但如果你用了同步的写法，后果可能不是报错，而是静默的性能退化：

# 错误写法 —— 能跑，但是没有异步效果
async def fetch_data():
    with aiohttp.ClientSession() as session:  # 忘了写 async
        async with session.get(url) as resp:
            return await resp.json()

async with 和 with 在 Python 里是两个完全不同的协议。前者调用 __aenter__ 和 __aexit__，后者调用 __enter__ 和 __exit__。如果你对异步上下文管理器用了普通 with，Python 会用同步协议去调用，很多情况下不会报错，因为对象可能恰好也实现了同步协议，但资源管理可能就没走异步路径了。

自己实现异步上下文管理器的正确姿势：

class AsyncResource:
    async def __aenter__(self):
        self.conn = await create_async_connection()
        return self.conn

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

也可以用 contextlib.asynccontextmanager：

from contextlib import asynccontextmanager

@asynccontextmanager
async def async_resource():
    conn = await create_async_connection()
    try:
        yield conn
    finally:
        await conn.close()

这里的规则和同步版一样：yield 必须在 try 里，清理在 finally 里。

五：上下文管理器和装饰器混用时的作用域问题

有时候你会想写一个既是上下文管理器又是装饰器的东西。比如一个临时修改环境变量的工具：

import os
from contextlib import contextmanager

@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:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v

这个 temp_env 只能当上下文管理器用，不能当装饰器。如果你想让一个函数变成既能 with temp_env(...) 又能 @temp_env(...)，需要这样写：

from contextlib import ContextDecorator

class temp_env(ContextDecorator):
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.old = {}

    def __enter__(self):
        self.old = {k: os.environ.get(k) for k in self.kwargs}
        os.environ.update(self.kwargs)
        return self

    def __exit__(self, *args):
        for k, v in self.old.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v
        return False

继承 ContextDecorator 之后，这个类就可以直接当装饰器用了：

@temp_env(DEBUG='true', LOG_LEVEL='debug')
def run_tests():
    # 这个函数执行期间，DEBUG 和 LOG_LEVEL 会被临时修改
    assert os.environ['DEBUG'] == 'true'

run_tests()
# 函数执行完，环境变量自动恢复

# 也可以当上下文管理器用
with temp_env(DEBUG='true'):
    ...

ContextDecorator 是标准库自带的，不需要额外安装。很多 Python 开发者不知道它的存在，但它很适合写测试辅助工具。

一个实际场景：数据库事务的自动回滚

最后分享一个我在项目里真实使用的上下文管理器——处理数据库事务：

class Transaction:
    def __init__(self, conn):
        self.conn = conn

    def __enter__(self):
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            self.conn.rollback()
            logger.error(f"事务回滚: {exc_val}")
        else:
            self.conn.commit()
        return False  # 不吞异常

# 使用
with Transaction(db) as conn:
    conn.execute("UPDATE accounts SET balance = balance - 100 WHERE id = 1")
    conn.execute("UPDATE accounts SET balance = balance + 100 WHERE id = 2")
    # 如果中间抛异常，自动回滚

这里面有两个关键点：return False 保证异常不会被吞掉（用户需要知道出错了），回滚操作只记录日志不往上抛新的异常（否则会覆盖原始异常）。

最后

上下文管理器是 Python 里一个设计得很好的特性，但用好不容易。回想一下，我见过的大部分资源泄漏问题，根因都不是什么高深的并发 bug，就是上下文管理器没写对。

如果你的项目里也有自定义的上下文管理器，不妨回去检查一下：__exit__ 有没有不小心返回 True、@contextmanager 里的 yield 有没有包在 try 里。花五分钟改一下，可能省下一整天的排查时间。