Python上下文管理(Context Manager)与contextlib快速入门

1. 为什么要关注资源管理？

在Python日常开发中，文件、数据库连接、网络请求、临时目录这类需要用完清理的资源随处可见。如果只写“打开/申请资源 → 使用”的代码，一旦中间报错，清理步骤就会被跳过，导致内存泄漏、连接池耗尽、临时文件残留等问题。

最原始的修复方法是用try...finally硬绑定：

try:
    f = open('/path/to/demo.txt', 'r', encoding='utf-8')
    data = f.read()
finally:
    # 即使前面read()出错，也会走到这里
    if f:
        f.close()

但这种写法太啰嗦，清理逻辑和业务逻辑混在一起，还容易忘记加if f这类边界判断。

2. Python原生解决方案：with语句

Python 2.5引入的with就是专门解决这类问题的语法糖，它自动处理“申请/进入 → 使用 → 清理/退出”的流程，哪怕中间抛异常也不例外：

with open('/path/to/demo.txt', 'r', encoding='utf-8') as f:
    data = f.read()

短短两行，把资源绑定和清理都搞定了，干净又安全。

3. 什么对象能放进with？上下文协议！

不是所有东西都能跟with玩，只有实现了上下文管理协议的对象才行——也就是要有两个魔法方法：

方法触发时机返回值作用
__enter__进入with代码块前赋给as后面的变量（可选）
__exit__离开with代码块时（无论对错）返回True可阻止异常向上抛

3.1 手写一个数据库连接上下文类

我们可以先写一个最基础的实现感受一下：

class DatabaseConnection:
    def __init__(self, db_name: str):
        self.db_name = db_name
        self.connection = None

    def __enter__(self) -> "DatabaseConnection":
        print(f"🔗 正在连接数据库 {self.db_name}")
        self.connection = f"MockConnection({self.db_name})"
        return self  # 把自己返回给with后的变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"🔒 正在关闭数据库连接 {self.db_name}")
        self.connection = None
        # 这三个参数是系统传的：类型、值、堆栈追踪，无异常时都是None
        if exc_type:
            print(f"⚠️ 捕获到异常：{exc_type.__name__} - {exc_val}")
        return True  # 比如不想让数据库查询的小错误打断整个程序？可以这样用

    def query(self, sql: str):
        print(f"📝 执行SQL：{sql}")
        return [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]

# 实际使用
with DatabaseConnection("local_test") as db:
    users = db.query("SELECT * FROM users")
    print(f"📦 查询结果：{users}")

运行后会发现，即使query()里抛异常，🔒 关闭连接的日志也一定会打印。

4. 用contextlib偷懒：避免手写魔法方法！

手写魔法方法虽然清晰，但对于简单场景还是太麻烦——比如临时目录、加锁这种只需要“进入做一件事，退出做另一件事”的。Python标准库的contextlib就是专门用来简化上下文管理器实现的。

4.1 最常用：@contextmanager装饰器

这个装饰器能把带yield的生成器函数直接变成上下文管理器，把原本__enter__/__exit__的逻辑拆成三段：

  1. yield前的代码 → 对应__enter__
  2. yield出来的东西 → 对应__enter__的返回值
  3. yield后的代码 → 对应__exit__（记得包在try...finally里！）

重写数据库连接

from contextlib import contextmanager

@contextmanager
def database_connection(db_name: str):
    print(f"🔗 正在连接数据库 {db_name}")
    connection = f"MockConnection({db_name})"
    db_conn_obj = type("", (), {"connection": connection, "query": lambda self, sql: print(f"📝 {sql}") or []})()
    try:
        yield db_conn_obj  # 中间的业务逻辑在这里执行
    finally:
        print(f"🔒 正在关闭数据库连接 {self.db_name}")

更简单的例子：HTML标签嵌套

这个例子不需要返回值，用来做代码块的包裹刚好：

@contextmanager
def html_tag(tag_name: str):
    print(f"<{tag_name}>")
    yield
    print(f"</{tag_name}>")

# 嵌套使用也很方便
with html_tag("div"):
    with html_tag("h1"):
        print("Hello Python Context Manager!")
    with html_tag("p"):
        print("这是用@contextmanager生成的标签上下文。")

运行输出：

<div>
<h1>
Hello Python Context Manager!
</h1>
<p>
这是用@contextmanager生成的标签上下文。
</p>
</div>

4.2 适配器工具：closing()

有些第三方库的对象有close()方法，但没有实现上下文协议（比如旧版urllib的urlopen），这时候可以用closing()给它“套个壳”：

from contextlib import closing
from urllib.request import urlopen

# urlopen返回的对象没有__enter__/__exit__，但有close()
with closing(urlopen("https://www.python.org")) as page:
    # 只打印前2行的前100个字符
    for i, line in enumerate(page):
        if i >= 2:
            break
        print(line.decode("utf-8")[:100])

4.3 批量管理工具：ExitStack

如果需要同时打开数量不确定的多个上下文（比如动态传入的文件名列表），用嵌套with根本不可能，这时候可以用ExitStack：

from contextlib import ExitStack
from pathlib import Path

def merge_files(input_files: list[Path], output_file: Path):
    with ExitStack() as stack:
        # 动态把所有输入文件加入栈管理
        files = [stack.enter_context(open(f, "r", encoding="utf-8")) for f in input_files]
        # 同时把输出文件也加入
        out = stack.enter_context(open(output_file, "w", encoding="utf-8"))
        
        # 业务逻辑：合并所有文件
        for f in files:
            out.write(f.read())
            out.write("\n=== 分割线 ===\n")

# 使用示例
merge_files(
    input_files=[Path("file1.txt"), Path("file2.txt"), Path("file3.txt")],
    output_file=Path("merged.txt")
)

ExitStack会在退出时按加入的反顺序自动清理所有资源。

5. 现代Python中的最佳实践

  1. 优先用with：只要是需要清理的资源（文件、锁、数据库连接、信号量），都先想能不能用with
  2. 优先用contextlib简化：简单场景用@contextmanager，适配器用closing()，批量用ExitStack
  3. @contextmanager记得加try...finally：否则业务逻辑抛异常的话，yield后的清理代码不会执行
  4. __exit__别随便返回True：除非你确定要吞掉这个异常，否则尽量让它正常向上抛
  5. 类型标注要跟上：Python 3.6+可以用typing.ContextManager（或者Python 3.9+的contextlib.AbstractContextManager）标注生成器函数的返回值，提升可读性

6. 两个超实用的实际例子

6.1 自动删除的临时目录

用标准库的tempfile.mkdtemp创建临时目录后，很容易忘记删，用@contextmanager包一下就好：

import tempfile
import shutil
from pathlib import Path
from contextlib import contextmanager

@contextmanager
def auto_clean_temp_dir(prefix: str = "tmp_", suffix: str = "") -> Path:
    temp_dir = Path(tempfile.mkdtemp(prefix=prefix, suffix=suffix))
    try:
        yield temp_dir
    finally:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
            print(f"🗑️ 已删除临时目录：{temp_dir}")

# 使用示例
with auto_clean_temp_dir(prefix="test_context_") as td:
    # 在临时目录里随便写文件
    test_file = td / "test.txt"
    test_file.write_text("这是临时文件的内容")
    print(f"📄 已在临时目录创建文件：{test_file}")

6.2 数据库事务自动回滚/提交

这个在Web开发（比如Flask-SQLAlchemy、FastAPI-SQLAlchemy）里非常常见：

# 假设session是SQLAlchemy的会话对象
from contextlib import contextmanager

@contextmanager
def auto_transaction(session):
    # 如果已经在事务里，就不嵌套，直接yield
    if session.in_transaction():
        yield
        return
    try:
        with session.begin():
            yield  # session.begin()已经包含了commit/rollback逻辑，但我们可以加个日志
        print("✅ 事务已提交")
    except Exception as e:
        print(f"❌ 事务已回滚：{e}")
        raise  # 别吞掉异常，让上层处理

通过掌握上下文管理器和contextlib，你就能写出更Pythonic、更安全、更简洁的资源管理代码啦！