Python 的 ExitStack 简直是资源管理神器，但你可能还没 get 到

如果你写过需要同时打开多个文件、多个锁或者多个数据库连接的代码，肯定见过这种场景：

f1 = open('a.txt', 'r')
f2 = open('b.txt', 'r')
try:
    data = f1.read() + f2.read()
finally:
    f1.close()
    f2.close()

如果还要处理异常，再嵌套一层 with，代码瞬间就变成“箭头形”。而且一旦文件数量变成动态的，比如列表里有几十个文件路径，上面的写法基本就没法看了。这时候 contextlib.ExitStack 就能派上大用场。

ExitStack 是什么？

简单说，它是一个可以帮助你动态管理任意多个上下文管理器的工具。它本身也是一个上下文管理器，但它的内部有一个“栈”，你可以往里面压入各种上下文管理器。当 with ExitStack() 块结束时，栈里的所有上下文管理器会按照后进先出的顺序自动执行清理操作，相当于帮你把所有的 finally 放在了一起。

更妙的是，你可以根据运行条件决定到底要不要把某个资源注册进去。常见场景：

• 数量不确定的文件打开
• 按需获取的锁
• 临时环境变量设置
• 动态添加的自定义清理回调

第一个例子：动态打开多个文件

假设我现在要合并一堆文本文件的内容，文件数量不定，但所有文件都要在最后关闭。用 ExitStack 可以这样写：

import contextlib

file_paths = ['a.txt', 'b.txt', 'c.txt']

with contextlib.ExitStack() as stack:
    files = []
    for path in file_paths:
        f = open(path, 'r', encoding='utf-8')
        files.append(f)
        stack.callback(f.close)  # 注册关闭函数

    contents = [f.read() for f in files]

# 出了 with 块，所有文件已经自动关闭
print(contents)

这里用 stack.callback(f.close) 把关闭函数压入栈，而不是用 stack.enter_context(f)。两者有什么区别？

enter_context() 会调用上下文管理器的 __enter__ 方法，然后返回其 __exit__ 方法需要的对象。对于文件对象，使用 enter_context(f) 和直接 with f: 是一样的。而 callback(f.close) 则是把 f.close 当作一个普通的清理函数，ExitStack 在退出时会自动调用它，不管是否发生异常。

用 callback 的好处是你不用非得是上下文管理器，任何可调用对象都能注册，比如打印日志、释放临时变量。

我个人更喜欢 enter_context，因为 __exit__ 还能接收异常信息，方便处理错误。上面的文件读取场景，用 enter_context 也很清晰：

with contextlib.ExitStack() as stack:
    files = [stack.enter_context(open(path, 'r', encoding='utf-8')) for path in file_paths]
    contents = [f.read() for f in files]
# 自动关闭

注意：这里 open() 返回的文件对象本身就是上下文管理器，所以 enter_context 会调用它的 __enter__ 并返回文件对象本身。当 with 块退出时，每个文件都会自动 close。

第二个案例：临时变更环境变量

有些脚本需要临时改变环境变量，比如给子进程设置一个自定义的 PYTHONPATH。传统写法是存旧值，改新值，最后恢复。如果是异常中断，还得保证恢复逻辑能执行。用 ExitStack 的 callback 可以优雅搞定：

import contextlib
import os

def set_env_var(name, value):
    old_value = os.environ.get(name)
    os.environ[name] = value

    if old_value is None:
        # 如果之前没有这个变量，那清理时就把它删掉
        return lambda: os.environ.pop(name, None)
    else:
        # 否则恢复原值
        return lambda: os.environ.__setitem__(name, old_value)

with contextlib.ExitStack() as stack:
    restore1 = set_env_var('API_KEY', 'abc123')
    restore2 = set_env_var('DEBUG', 'true')

    stack.callback(restore1)
    stack.callback(restore2)

    # 在这里环境变量就是临时的
    print(os.environ['API_KEY'], os.environ['DEBUG'])

# 退出后环境变量自动恢复
print(os.environ.get('API_KEY'))  # 原值

注意 set_env_var 返回的是一个闭包，这个闭包就是清理函数。把它们都压入栈后，即使中间抛异常，所有回调仍然会执行。相比手动管理 try/finally 不知道高到哪里去了。

第三个案例：按需获取多个锁

多线程编程里锁的获取顺序很重要，但也要保证所有锁都能被释放。假设有两把锁 lock_a 和 lock_b，要根据条件决定是否获取，而且获取后一定要释放：

import contextlib
import threading

lock_a = threading.Lock()
lock_b = threading.Lock()

def critical_section(need_a=True, need_b=True):
    with contextlib.ExitStack() as stack:
        if need_a:
            stack.enter_context(lock_a)
        if need_b:
            stack.enter_context(lock_b)

        # 同时持有两把锁，执行关键操作
        print("执行中...")
    
    # 锁已自动释放

这里 enter_context 接收锁对象，因为它实现了上下文管理器协议。当 with 块结束时，无论是正常结束还是抛异常，锁都会被正确释放。这比手动写 acquire/release 要安全得多。

如果你需要同时等待多个条件变量（比如生产者消费者），ExitStack 也可以派上用场。

更高级的用法：临时目录与 chdir

写测试或脚本时经常要临时切换工作目录，然后在结束后恢复。用 ExitStack 就能这样：

import contextlib
import os
import tempfile

with contextlib.ExitStack() as stack:
    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    stack.callback(lambda: os.rmdir(temp_dir))

    # 保存当前目录，然后切换到临时目录
    old_cwd = os.getcwd()
    os.chdir(temp_dir)
    stack.callback(lambda: os.chdir(old_cwd))

    print("当前工作目录:", os.getcwd())
    # 做一些操作...

# 退出后自动恢复原来的目录，并删除临时目录（如果为空的话）
print("恢复后的目录:", os.getcwd())

这里面有几个坑需要注意：os.rmdir 只能删空目录，如果你在临时目录里创建了文件，需要先删文件。所以更稳妥的方式是使用 tempfile.TemporaryDirectory，但这里主要是演示 callback 的组合能力。

实际上，ExitStack 还提供了 pop_all() 方法，可以把已注册的清理回调全部弹出并返回一个新的 ExitStack 对象。这样你可以把清理责任转移给另一个栈，实现“延迟清理”或者“部分清理”的效果。这个特性在处理一些需要手动提交/回滚的事务时非常有用。

事务中的提交或回滚：一个真实的业务案例

在做数据库操作的时候，如果多个操作里有一个失败，就需要整体回滚。用 ExitStack 可以做到：先记录所有要执行的清理操作，然后根据事务结果决定是提交还是回滚。

下面是一个简化版：

import contextlib

# 假设这是一个数据库连接对象
class Database:
    def __init__(self):
        self.log = []

    def perform_action(self, action):
        self.log.append(action)
        print(f"执行操作: {action}")

    def commit(self):
        self.log.append("commit")
        print("提交事务")

    def rollback(self):
        self.log.append("rollback")
        print("回滚事务")

def process_transaction(db, actions):
    with contextlib.ExitStack() as stack:
        # 把回滚操作注册为回调，但先不执行
        stack.callback(db.rollback)
        stack.callback(lambda: print("事务结束"))

        for action in actions:
            if action == "bad":
                raise ValueError("遇到错误，需要回滚")
            db.perform_action(action)

        # 全部成功，则取消回滚，改为提交
        # pop_all() 会取出当前所有的回调，同时从当前栈中移除
        # 然后我们再把提交操作压入新栈
        pending_stack = stack.pop_all()
        db.commit()
        pending_stack.close()  # 这会执行里面对应的清理，但此时已经没有rollback了

这段代码的思路是：一开始把 db.rollback 注册为回调，如果操作过程中抛出异常，ExitStack 会自动调用 rollback。如果所有操作成功，我们调用 pop_all() 把原来的回滚回调全部取出，再手动 commit()。因为回调已经被弹出，所以新的栈里没有回滚操作，此时调用 close() 就只执行剩下的清理函数。

虽然这个例子有点抽象，但《Python Cookbook》里也有类似技巧。它很好地展示了 ExitStack 不仅能做“物理资源清理”，还能做“逻辑回滚”。

ExitStack 使用的一些小技巧

• enter_context 会返回上下文管理器的 __enter__ 返回值，你可以保存这个返回值，方便后续使用。
• callback 可以接受任意参数和关键字参数，比如 stack.callback(os.remove, tmp_file)。
• 多个 callback 的调用顺序是后进先出（LIFO）。如果你依赖清理的顺序，注意压栈的顺序。
• ExitStack 自身也可以作为其他 ExitStack 的上下文管理器，实现嵌套组合。

总结：你还需要手写 try/finally 吗？

现在我再看到那种有一堆 try 然后里面嵌套 with 再嵌套 finally 的代码，心里就在想：为什么不直接用 ExitStack？它把“进入资源”和“退出清理”完全解耦，让你的代码更线性，更容易阅读。

当然，它也不是银弹。对于简单的两个资源，直接用两个 with 也没问题。但一旦数量不固定，或者清理逻辑需要动态调整，ExitStack 就是最优雅的工具。

如果你之前没关注过它，那么下次再遇到资源管理的烦恼时候，不妨先想起 contextlib.ExitStack。它能帮你省掉一堆无聊的 finally 代码，还能避免遗漏关闭资源。

我在项目里就是用它重写了一个老旧的配置加载模块，不仅代码变短了，而且多个配置文件的异常处理也变得更统一。好的工具就是能让你越用越舒服。