Python 与数据实践 · 第 12 章

with 与上下文协议

写代码久了你会发现一件事：大量资源都讲究"有借有还"——打开的文件要关、连上的服务要断、占到的锁要放、申请的句柄要释放。麻烦在于，"还"这个动作太容易掉链子：一个提前的 return、一个半路杀出的异常，清理代码就被整个…

约 12 分钟4,794 字12 个代码示例

写代码久了你会发现一件事：大量资源都讲究"有借有还"——打开的文件要关、连上的服务要断、占到的锁要放、申请的句柄要释放。麻烦在于，"还"这个动作太容易掉链子：一个提前的 return、一个半路杀出的异常，清理代码就被整个跳过去了。with 语句就是 Python 对这件事的官方解法：把"借"和"还"写进一套协议，让解释器替你盯着收尾——不管中间发生什么，离开代码块时该做的清理一定会做。

1. with 的长相与内脏

语法长这样：

with 表达式 as 变量:
    代码块

这行语法糖脱掉外壳之后，大致等效于：

变量 = 表达式.__enter__()
try:
    代码块
except Exception as e:
    swallowed = 变量.__exit__(type(e), e, e.__traceback__)
    if not swallowed:
        raise
else:
    变量.__exit__(None, None, None)

从这段展开式里能直接读出几件事：

踏进代码块之前，先调 __enter__()，它的返回值交给 as 后面的变量
离开代码块时，无论正常走完还是中途炸了，__exit__() 都会被调用——自动清理的底气就在这里
出异常时，__exit__ 会收到三个参数：异常类型、异常对象、调用栈信息；没出异常，三者全是 None
__exit__ 的返回值是个开关：返回 True 表示"这个异常归我处理"，异常就地消亡；返回 False（或者不写 return），异常继续向上传播

2. 上下文管理器：一纸双方法协议

同时实现了 __enter__ 和 __exit__ 两个方法的对象，就叫上下文管理器（Context Manager）。 with 语句只认这种对象——两个方法缺了任何一个，with 都会直接报错。

3. 最经典的搭档：文件

# 文件操作 —— 关闭动作全自动
with open("notes.txt", "r") as f:
    text = f.read()
    # 走出 with 块的那一刻，文件已被关好

# 不借 with 之手的等价写法（体会一下手动管理的啰嗦）：
f = open("notes.txt", "r")
try:
    text = f.read()
except Exception as e:
    swallowed = f.__exit__(type(e), e, e.__traceback__)
    if not swallowed:
        raise
else:
    f.__exit__(None, None, None)

4. 亲手造一个上下文管理器

协议就两个方法，照着填就行。做一个消息服务器的连接：

class MessageBroker:
    def __init__(self, server):
        self.server = server
        self.online = False

    def __enter__(self):
        """踏进 with 块时执行；返回值赋给 as 后面的变量"""
        print(f"已连上消息服务器: {self.server}")
        self.online = True
        return self  # 把自己交出去，块内就能用 conn.xxx

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        踏出 with 块时执行（出不出异常都会到这一步）
        exc_type / exc_val / exc_tb：出异常时是"异常三件套"，没出全是 None
        返回 True = 异常到此为止；返回 False = 异常照常上报
        """
        print(f"已断开消息服务器: {self.server}")
        self.online = False
        return False  # 不拦异常，出了事让上层知道

    def publish(self, topic, message):
        if not self.online:
            raise RuntimeError("尚未连接消息服务器")
        print(f"向频道 {topic} 发送: {message}")


with MessageBroker("mq.internal") as conn:
    print(f"在线状态: {conn.online}")          # True
    conn.publish("订单频道", "新订单已创建")

# 走出 with 块之后
print(f"在线状态: {conn.online}")              # False

# 输出：
# 已连上消息服务器: mq.internal
# 在线状态: True
# 向频道 订单频道 发送: 新订单已创建
# 已断开消息服务器: mq.internal
# 在线状态: False

5. `exit` 的返回值：异常的拦截阀

__exit__ 不只负责收拾残局，还可以顺手处理甚至"吞掉"异常：

class ErrorMuffler:
    """名单之内的异常一律消音，名单之外的照常放行"""

    def __init__(self, *muted_types):
        self.muted_types = muted_types

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None and issubclass(exc_type, self.muted_types):
            print(f"已消音: {exc_type.__name__}: {exc_val}")
            return True   # 返回 True：异常已被处理，外界毫无感知
        return False      # 不在名单里的异常，不归我管


with ErrorMuffler(ZeroDivisionError):
    result = 1 / 0        # 不会崩
    print("到不了这一行")

print("程序若无其事地继续")  # 正常执行

6. 借生成器之力：@contextmanager 装饰器

第 11 章刚学完生成器，这里立刻派上用场。写类有时候显得太隆重，contextlib 模块提供了 @contextmanager 装饰器——用一个带 yield 的函数就能表达同样的协议：

from contextlib import contextmanager


@contextmanager
def camera(session):
    """yield 之前的代码 ≈ __enter__，yield 之后的代码 ≈ __exit__"""
    print(f"开启摄像头: {session}")
    handle = {"device": session, "power": "on"}
    try:
        yield handle  # 暂停点：with 块在这里执行，handle 交给 as 变量
    finally:
        print(f"关闭摄像头: {session}")  # 块内出不出事，这行都会执行


with camera("前台闸机") as cam:
    print(f"抓拍中: {cam}")
# 开启摄像头: 前台闸机
# 抓拍中: {'device': '前台闸机', 'power': 'on'}
# 关闭摄像头: 前台闸机

再看一个带完整异常处理的版本——先写临时文件，确认成功才替换正式文件，失败就清场：

from contextlib import contextmanager


@contextmanager
def atomic_write(path):
    """先把内容写进临时文件，确认无误再顶替正式文件"""
    temp_path = path + ".swap"
    f = open(temp_path, "w")
    try:
        yield f
        f.close()
        # 全部写完且没出岔子，正式换名
        import os
        os.replace(temp_path, path)
        print("落盘成功")
    except Exception as e:
        # 中途出错：残缺的临时文件不要了
        f.close()
        import os
        if os.path.exists(temp_path):
            os.remove(temp_path)
        print(f"落盘失败: {e}")
        raise  # 异常继续抛出去，交给上层定夺


with atomic_write("settings.json") as f:
    f.write('{"theme": "dark"}\n')

7. 一条 with 掌管多个资源

资源一多，with 可以层层嵌套，也可以写成一行：

# 嵌套写法
with open("原始日志.txt", "r") as src:
    with open("归档副本.txt", "w") as backup:
        backup.write(src.read().upper())

# 逗号并列写法（Python 3.1+），效果与上面完全相同
with open("原始日志.txt", "r") as src, open("归档副本.txt", "w") as backup:
    backup.write(src.read().upper())

退出顺序与进入顺序相反：后进入的先退出——backup 先关，src 后关。

8. 作业

动笔之前，先把 __enter__ 和 __exit__ 的调用时机、参数含义、返回值含义复述一遍，复述顺畅再开工。

一、给代码块计时

用 @contextmanager 实现一个计时上下文管理器 elapsed：

from contextlib import contextmanager

# 目标效果：
with elapsed("数据清洗"):
    import time
    time.sleep(1)
    print("清洗完成")
# 清洗完成
# 数据清洗 耗时: 1.0013 秒

二、临时工作区

编写一个 ScratchSpace 上下文管理器：进入时开辟一个临时目录，退出时连同里面所有内容一并铲除：

with ScratchSpace() as work_dir:
    print(f"工作区: {work_dir}")
    # 可以在目录里随便建文件、写东西
    # 离开 with 块后，目录连同内容自动消失

print("工作区已清理")

提示： 用 tempfile 模块创建临时目录，用 shutil.rmtree 连根删除。

三、辨析题：with 能实现"自动重试"吗？

一个诱人的想法：写个上下文管理器，块内代码抛异常就自动再执行一遍，直到成功或达到次数上限。先别急着动手——想想 with 的进出机制，允许同一个代码块被"重复进入"吗？写下你的结论和理由，再对照 answers/p3.py（不可行方案的剖析）与 answers/p3-1.py（用第 10 章的装饰器给出的正解）。

四、思考题

下面代码的输出是什么？为什么？

from contextlib import contextmanager


@contextmanager
def stage_lights():
    print("灯光就位")
    yield
    print("落幕关灯")


with stage_lights():
    print("演出进行中")
    raise RuntimeError("道具故障")
    print("这行永远轮不到")

如果改成下面的写法，输出又有什么不同？

@contextmanager
def stage_lights():
    print("灯光就位")
    try:
        yield
    except Exception as e:
        print(f"临时救场: {e}")
    finally:
        print("落幕关灯")


with stage_lights():
    print("演出进行中")
    raise RuntimeError("道具故障")

下一章：13 · 异常体系与处理