Python 开发技巧 · Python 上下文管理器 —— 从 with 到 contextlib 实战

1. 知识点简介

上下文管理器是 Python 中最被低估的「资源管理利器」。with 语句背后是 __enter__ 和 __exit__ 协议，确保资源无论是否发生异常都能正确释放。标准库 contextlib 更进一步，提供了装饰器、工具函数来简化上下文管理器的编写。

适用场景：

• 文件/网络连接/数据库事务的自动关闭
• 临时修改全局状态（环境变量、目录切换、配置覆写）
• 性能计时、日志上下文、锁获取释放
• 异常吞没与转换

---

2. 基础：实现上下文管理器的三种方式

2.1 基于类（__enter__ / __exit__）

class Timer:
    """计时器上下文管理器"""
    def __enter__(self):
        self.start = time.perf_counter()
        return self  # as 子句拿到的对象

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"⏱️ 耗时: {self.elapsed:.4f}s")
        return False  # False = 不吞没异常，True = 吞没

# 使用
import time
with Timer() as t:
    sum(range(10_000_000))
# 输出: ⏱️ 耗时: 0.2845s

__exit__ 返回值的含义：

• False（或 None）：异常继续传播
• True：异常被吞没，with 块结束后继续正常执行

2.2 基于 @contextmanager 装饰器（推荐）

from contextlib import contextmanager

@contextmanager
def timed(description: str = ""):
    """用生成器语法实现上下文管理器"""
    start = time.perf_counter()
    try:
        yield  # ← 这里是 __enter__ 和 __exit__ 的分界线
    finally:
        elapsed = time.perf_counter() - start
        print(f"⏱️ {description}: {elapsed:.4f}s")

# 使用
with timed("大量计算"):
    sum(range(20_000_000))

原理 ：yield 之前 = __enter__，yield 之后 = __exit__。用 try/finally 确保无论如何都会执行清理。

2.3 基于 contextlib.ContextDecorator（类+装饰器混合）

from contextlib import ContextDecorator

class timed(ContextDecorator):
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, *exc):
        elapsed = time.perf_counter() - self.start
        print(f"⏱️ 耗时: {elapsed:.4f}s")
        return False

# 既可以当装饰器
@timed()
def heavy_work():
    sum(range(10_000_000))

heavy_work()

# 也可以当上下文管理器
with timed():
    sum(range(10_000_000))

---

3. contextlib 工具箱

3.1 contextlib.suppress ------ 优雅忽略异常

import os
from contextlib import suppress

# ❌ 传统写法
try:
    os.remove("temp.txt")
except FileNotFoundError:
    pass

# ✅ suppress 写法
with suppress(FileNotFoundError):
    os.remove("temp.txt")

# 支持多个异常类型
with suppress(FileNotFoundError, PermissionError):
    os.remove("/protected/temp.txt")

3.2 contextlib.redirect_stdout / redirect_stderr ------ 临时重定向输出

from contextlib import redirect_stdout, redirect_stderr
import io

buf = io.StringIO()
with redirect_stdout(buf), redirect_stderr(buf):
    print("这段文字被捕获了")
    print("错误信息也进去了")

output = buf.getvalue()
print(f"捕获内容: {output!r}")

3.3 contextlib.closing ------ 给没有上下文协议的对象加上 with 支持

from contextlib import closing
import urllib.request

# urllib.request.urlopen 本身支持 with，但某些旧 API 不支持
with closing(urllib.request.urlopen("https://httpbin.org/get")) as resp:
    data = resp.read()
    print(f"状态码: {resp.status}")
# closing 确保 .close() 被调用

3.4 contextlib.ExitStack ------ 管理动态数量的上下文管理器

最有用的工具之一，适合数量不固定的资源管理场景。

from contextlib import ExitStack

def open_multiple_files(filenames: list[str]):
    """动态打开任意数量的文件"""
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fname))
            for fname in filenames
        ]
        # 所有文件在这里可用；ExitStack 退出时自动关闭全部
        return files  # ⚠️ 注意：这会导致文件被提前关闭！

⚠️ ExitStack 的经典坑 ：如果 ExitStack.__exit__ 在 return 之前执行，文件已经被关闭。修正方式：

def read_all_files(filenames: list[str]) -> list[str]:
    """正确做法：在 with 块内完成全部读写操作"""
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fname))
            for fname in filenames
        ]
        return [f.read() for f in files]  # ✅ 在退出前完成操作

更酷的用法：延迟回调注册

from contextlib import ExitStack

with ExitStack() as stack:
    stack.callback(lambda: print("清理 1"))
    stack.callback(lambda: print("清理 2"))
    # 退出时按 LIFO 顺序执行回调
    # 输出: 清理 2 → 清理 1

---

4. 实战：业务场景上下文管理器

4.1 数据库事务自动回滚

@contextmanager
def transaction(cursor):
    print("🔵 开始事务")
    try:
        yield cursor
        cursor.execute("COMMIT")
        print("🟢 事务提交")
    except Exception:
        cursor.execute("ROLLBACK")
        print("🔴 事务回滚")
        raise  # 重新抛出异常

# 使用
import sqlite3
conn = sqlite3.connect(":memory:")
cur = conn.cursor()
cur.execute("CREATE TABLE test (id INT, name TEXT)")

with transaction(cur):
    cur.execute("INSERT INTO test VALUES (1, 'Alice')")
    # 如果这里发生异常，自动回滚

4.2 临时切换工作目录

import os
from contextlib import contextmanager

@contextmanager
def cd(path: str):
    old_cwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old_cwd)

# 使用
with cd("/tmp"):
    print(f"当前目录: {os.getcwd()}")  # /tmp
print(f"恢复目录: {os.getcwd()}")      # 原来的目录

4.3 临时修改环境变量

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

# 使用
with set_env(DEBUG="1", API_KEY="test-key"):
    print(os.environ["DEBUG"])  # 1
print(os.environ.get("DEBUG"))  # None（恢复）

4.4 连接池中的 Redis Pipeline

@contextmanager
def pipeline(redis_client):
    pipe = redis_client.pipeline()
    try:
        yield pipe
        pipe.execute()  # 批量执行
    except Exception:
        pipe.reset()    # 出错时清空
        raise

# 使用
with pipeline(redis_client) as pipe:
    pipe.incr("counter")
    pipe.set("key", "value")
    pipe.lpush("queue", "item")
# 3 条命令一次网络往返发送

---

5. 避坑指南

坑点	说明	解决
❌ @contextmanager 忘记 yield	生成器不 yield 会抛 RuntimeError	确保生成器有且仅有一个 yield
❌ __exit__ 返回 True 吞异常	无意识地让异常静默消失	只有明确想吞异常时才 return True
❌ ExitStack 返回内部资源	退出栈时资源已释放	在 with 块内用完资源再退出
❌ 嵌套 with 过于冗长	5-6 层嵌套可读性差	用 ExitStack 或逗号合并 with A(), B():
❌ 忽略 __exit__ 的参数	不处理 exc_type/exc_val/exc_tb 但没意识到	需要时用 issubclass(exc_type, SomeError) 判断

---

6. 总结

• 三选一 ：简单场景用 @contextmanager，类场景用 __enter__/__exit__，装饰器+上下文二合一用 ContextDecorator
• contextlib 四大金刚 ：suppress（吞异常）、redirect_stdout（重定向）、closing（补协议）、ExitStack（动态管理）
• 业务落地：事务、目录切换、环境变量、连接池 ------ 上下文管理器能帮你写出更安全的代码
• 核心原则 ：yield 是分界线，前面的初始化，后面的清理，try/finally 保底

用 with 不仅是「省了几行 close()」，而是在代码层面显式声明了资源的生命周期。好的代码告诉你「这是资源的边界」，而不是「记得自己关」。