Python 上下文管理器：with 语句的高级用法

Python 的 with 语句常用于自动管理资源，比如文件读写后自动关闭。但它的能力远不止于此。通过自定义上下文管理器，你可以控制任意代码块的进入和退出行为，实现更安全、更简洁的逻辑封装。

什么是上下文管理器？

上下文管理器是一个实现了 __enter__ 和 __exit__ 方法的对象。当你用 with 语句包裹一段代码时，Python 会在进入代码块前调用 __enter__，在离开时（无论是否出错）调用 __exit__。

创建一个最简单的上下文管理器：

class MyContext:
    def __enter__(self):
        print("进入上下文")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出上下文")
        return False  # 不抑制异常

使用它：

with MyContext() as ctx:
    print("执行中")

输出：

进入上下文
执行中
退出上下文

使用 contextlib 简化编写

对于简单场景，不必定义完整类。使用 contextlib.contextmanager 装饰器，把函数变成上下文管理器：

from contextlib import contextmanager

@contextmanager
def my_context():
    print("准备资源")
    try:
        yield "资源对象"
    finally:
        print("清理资源")

使用方式相同：

with my_context() as res:
    print(f"使用 {res}")

输出：

准备资源
使用 资源对象
清理资源

注意：yield 之前的代码相当于 __enter__，之后的 finally 块相当于 __exit__。即使发生异常，finally 也会执行。

高级技巧：处理异常

__exit__ 方法接收三个参数：exc_type、exc_val、exc_tb，分别代表异常类型、值和追踪信息。如果没异常，三者都是 None。

返回 True 可以抑制异常（不抛出）；返回 False 或 None 则让异常继续传播。

例如，实现一个“忽略特定异常”的上下文管理器：

from contextlib import contextmanager

@contextmanager
def ignore_exception(*exception_types):
    try:
        yield
    except exception_types:
        print("已忽略指定异常")
        pass  # 不重新抛出

使用：

with ignore_exception(ValueError, ZeroDivisionError):
    1 / 0  # 触发 ZeroDivisionError

print("程序继续运行")

输出：

已忽略指定异常
程序继续运行

嵌套与组合多个上下文

当需要同时管理多个资源时，可以嵌套 with 语句：

with open('input.txt') as f1, open('output.txt', 'w') as f2:
    f2.write(f1.read())

但若层级太深，可读性下降。使用 contextlib.ExitStack 动态组合多个上下文：

from contextlib import ExitStack

with ExitStack() as stack:
    files = [
        stack.enter_context(open(f'file{i}.txt'))
        for i in range(3)
    ]
    # 所有文件将在 with 结束时自动关闭
    for f in files:
        print(f.readline())

ExitStack 特别适合在运行时决定需要哪些资源（比如根据配置打开不同数量的文件或连接）。

实战案例：数据库事务管理

假设你有一个数据库连接对象 db，支持 commit() 和 rollback()。实现一个事务上下文管理器：

from contextlib import contextmanager

@contextmanager
def transaction(db):
    db.begin()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise  # 重新抛出异常

使用：

with transaction(db) as conn:
    conn.execute("INSERT INTO users ...")
    conn.execute("UPDATE accounts ...")
    # 如果中间出错，自动回滚；否则提交

这种方式确保了事务的原子性，且代码清晰无冗余。

自定义资源锁

多线程环境中，常需加锁。封装一个带超时和日志的锁上下文：

import threading
import time

class LoggedLock:
    def __init__(self, name, timeout=5):
        self.name = name
        self.timeout = timeout
        self.lock = threading.Lock()

    def __enter__(self):
        acquired = self.lock.acquire(timeout=self.timeout)
        if not acquired:
            raise TimeoutError(f"无法在 {self.timeout} 秒内获取锁 '{self.name}'")
        print(f"[{time.time():.2f}] 获取锁: {self.name}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.lock.release()
        print(f"[{time.time():.2f}] 释放锁: {self.name}")

使用：

lock = LoggedLock("my_lock")

with lock:
    time.sleep(1)
    # 模拟临界区操作

输出示例：

[1717020000.12] 获取锁: my_lock
[1717020001.12] 释放锁: my_lock

性能监控上下文

创建一个测量代码块执行时间的工具：

import time
from contextlib import contextmanager

@contextmanager
def timer(label="代码块"):
    start = time.perf_counter()
    try:
        yield
    finally:
        end = time.perf_counter()
        print(f"{label} 耗时: {end - start:.4f} 秒")

使用：

with timer("数据处理"):
    time.sleep(0.5)

输出：

数据处理 耗时: 0.5002 秒

可轻松扩展为记录日志、上报监控系统等。

注意事项与最佳实践

1. 始终在 __exit__ 或 finally 中清理资源，即使发生异常。
2. 不要在 __enter__ 中返回敏感对象，除非必要；通常返回 self 或一个代理对象。
3. 避免在 __exit__ 中吞掉所有异常（即返回 True），除非明确知道后果。
4. 优先使用 contextlib.contextmanager，除非需要复用状态或复杂逻辑。

测试上下文管理器的行为

验证你的上下文管理器是否正确处理异常：

def test_rollback_on_exception():
    class FakeDB:
        def __init__(self):
            self.committed = False
            self.rolled_back = False
        def commit(self): self.committed = True
        def rollback(self): self.rolled_back = True
        def begin(self): pass

    db = FakeDB()
    try:
        with transaction(db):
            raise ValueError("模拟错误")
    except ValueError:
        pass

    assert not db.committed
    assert db.rolled_back

运行此测试可确保事务逻辑可靠。

记住：任何需要“开始-结束”对称操作的地方，都值得考虑上下文管理器——无论是打开/关闭、加锁/解锁、启动/停止，还是记录/上报。它让资源管理变得自动化、防错且优雅。