Python 上下文管理器实战——with 语句、contextlib、异步上下文管理器

适读人群：想掌握 Python 资源管理最佳实践的工程师 | 阅读时长：约 14 分钟 | 核心价值：用上下文管理器优雅地管理任何资源，告别 try/finally 的重复代码

---

一次忘记关闭文件的教训

刚学 Python 的时候，我写了这样的代码：

f = open("data.txt", "r")
data = f.read()
# 忘记 f.close()

在 Java 里我会想到用 try-finally，但 Python 有更优雅的方式——with 语句。这个特性背后的机制叫"上下文管理器"，是 Python 资源管理的核心范式。

从文件到数据库连接，从锁到 HTTP 会话，从临时目录到测试 mock——只要涉及"获取-使用-释放"的模式，上下文管理器都是最好的解决方案。

这篇文章把上下文管理器的所有用法系统讲一遍，从基础到异步，从内置到自定义。

---

一、上下文管理器协议

with 语句背后是一个简单的协议：对象需要实现 __enter__ 和 __exit__ 方法。

class ManagedResource:
    def __enter__(self):
        print("资源获取")
        return self  # as 子句得到这个返回值

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("资源释放")
        # 返回 True 表示"我处理了这个异常，不要往上抛"
        # 返回 False 或 None 表示"让异常继续传播"
        if exc_type is not None:
            print(f"捕获到异常: {exc_type.__name__}: {exc_val}")
        return False  # 不吞掉异常

with ManagedResource() as r:
    print("使用资源")
    # 即使这里抛异常，__exit__ 也会被调用

和 Java try-with-resources 对比：

// Java
try (FileInputStream fis = new FileInputStream("file.txt")) {
    // 使用资源
}

# Python
with open("file.txt") as f:
    # 使用资源

---

二、contextlib：不写类，用生成器

实现上下文管理器不一定要写类，contextlib.contextmanager 让你用生成器函数代替：

from contextlib import contextmanager
import time

@contextmanager
def timer(name: str = ""):
    """计时上下文管理器"""
    start = time.perf_counter()
    try:
        yield  # with 块在这里执行
    finally:
        elapsed = (time.perf_counter() - start) * 1000
        print(f"[{name or 'timer'}] 耗时: {elapsed:.2f}ms")

with timer("数据库查询"):
    time.sleep(0.1)
    # 执行数据库操作

yield 前的代码是 __enter__，yield 后（finally 里）的代码是 __exit__。

---

三、实用上下文管理器集合

3.1 临时目录/文件

import tempfile
from pathlib import Path

# 临时目录（退出 with 后自动删除）
with tempfile.TemporaryDirectory() as tmpdir:
    tmp_path = Path(tmpdir) / "output.json"
    tmp_path.write_text('{"result": "ok"}')
    print(f"临时文件: {tmp_path}")
    # with 退出后，整个目录被删除

# 临时文件
with tempfile.NamedTemporaryFile(suffix=".csv", delete=True) as f:
    f.write(b"name,age\n")
    f.flush()
    print(f"临时文件路径: {f.name}")

3.2 改变工作目录

from contextlib import contextmanager
import os

@contextmanager
def chdir(path: str):
    """临时切换工作目录"""
    original = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(original)

with chdir("/tmp"):
    print(f"当前目录: {os.getcwd()}")
print(f"恢复目录: {os.getcwd()}")

3.3 数据库事务上下文

from contextlib import contextmanager

@contextmanager
def transaction(connection):
    """数据库事务管理器"""
    try:
        yield connection
        connection.commit()
        print("事务提交")
    except Exception as e:
        connection.rollback()
        print(f"事务回滚: {e}")
        raise

with transaction(db_connection) as conn:
    conn.execute("INSERT INTO users VALUES (?)", ("老张",))
    conn.execute("UPDATE stats SET count = count + 1")
    # 如果这里抛异常，会自动回滚

3.4 suppress：静默指定异常

from contextlib import suppress

# 相当于：
# try:
#     os.remove("file.txt")
# except FileNotFoundError:
#     pass
with suppress(FileNotFoundError):
    os.remove("file.txt")

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

3.5 redirect_stdout：重定向输出

from contextlib import redirect_stdout
import io

# 捕获 print 的输出
output = io.StringIO()
with redirect_stdout(output):
    print("这会被捕获，不会显示在终端")
    print("老张")

captured = output.getvalue()
print(f"捕获到的输出: {captured!r}")

---

四、异步上下文管理器

异步版本使用 __aenter__ 和 __aexit__，或者用 contextlib.asynccontextmanager：

from contextlib import asynccontextmanager
import asyncio

@asynccontextmanager
async def async_timer(name: str = ""):
    import time
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = (time.perf_counter() - start) * 1000
        print(f"[{name}] 异步耗时: {elapsed:.2f}ms")

async def main():
    async with async_timer("异步操作"):
        await asyncio.sleep(0.1)

---

五、完整可运行示例

#!/usr/bin/env python3
"""
Python 上下文管理器完整实战
"""
import asyncio
import os
import tempfile
import time
from contextlib import asynccontextmanager, contextmanager, suppress
from pathlib import Path
from typing import Any, Generator


# ===== 1. 计时器（yield 版）=====
@contextmanager
def timer(label: str = ""):
    start = time.perf_counter()
    try:
        yield
    finally:
        ms = (time.perf_counter() - start) * 1000
        print(f"⏱  [{label or 'elapsed'}]: {ms:.2f}ms")


# ===== 2. 临时环境变量 =====
@contextmanager
def env_vars(**kwargs: str) -> Generator[None, None, None]:
    """临时设置环境变量，退出后恢复原值"""
    original: dict[str, str | None] = {}
    for key, value in kwargs.items():
        original[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for key, old_value in original.items():
            if old_value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old_value


# ===== 3. 简单重试上下文 =====
@contextmanager
def retry_block(times: int = 3, delay: float = 0.0):
    """重试块：如果抛出异常，自动重试"""
    attempt = 0
    while True:
        try:
            yield attempt
            break
        except Exception as e:
            attempt += 1
            if attempt >= times:
                raise
            print(f"第{attempt}次失败，重试: {e}")
            if delay > 0:
                time.sleep(delay)


# ===== 4. 临时文件写入 =====
@contextmanager
def temp_csv(rows: list[list[str]]):
    """将数据写入临时 CSV，退出后删除"""
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".csv", delete=False, encoding="utf-8"
    ) as f:
        for row in rows:
            f.write(",".join(row) + "\n")
        path = f.name
    try:
        yield Path(path)
    finally:
        with suppress(FileNotFoundError):
            os.unlink(path)


# ===== 5. 类实现上下文管理器 =====
class MockDatabase:
    """模拟数据库连接"""

    def __init__(self, url: str):
        self.url = url
        self.connected = False
        self.operations: list[str] = []

    def __enter__(self) -> "MockDatabase":
        print(f"  连接数据库: {self.url}")
        self.connected = True
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"  关闭数据库连接 (异常: {exc_type.__name__ if exc_type else None})")
        self.connected = False
        return False  # 不吞掉异常

    def execute(self, sql: str) -> None:
        if not self.connected:
            raise RuntimeError("未连接到数据库")
        self.operations.append(sql)


# ===== 6. 异步上下文管理器 =====
@asynccontextmanager
async def managed_http_session():
    """管理 HTTP 会话生命周期（模拟）"""
    print("  创建 HTTP Session")
    session = {"active": True, "requests": 0}
    try:
        yield session
    finally:
        session["active"] = False
        print(f"  关闭 HTTP Session（共发出 {session['requests']} 个请求）")


# ===== 演示 =====
def demo_sync():
    print("=== 同步上下文管理器演示 ===\n")

    # 1. 计时器
    print("1. 计时器")
    with timer("排序操作"):
        data = sorted(range(100000), reverse=True)
    print()

    # 2. 环境变量
    print("2. 临时环境变量")
    print(f"  设置前 TEST_MODE = {os.environ.get('TEST_MODE', '(未设置)')}")
    with env_vars(TEST_MODE="true", LOG_LEVEL="DEBUG"):
        print(f"  设置中 TEST_MODE = {os.environ.get('TEST_MODE')}")
    print(f"  恢复后 TEST_MODE = {os.environ.get('TEST_MODE', '(未设置)')}")
    print()

    # 3. 临时文件
    print("3. 临时 CSV 文件")
    rows = [["name", "age", "score"], ["老张", "35", "95"], ["小陈", "28", "88"]]
    with temp_csv(rows) as csv_path:
        print(f"  临时文件: {csv_path}")
        print(f"  内容: {csv_path.read_text()!r}")
    print(f"  文件已删除: {not csv_path.exists()}")
    print()

    # 4. 数据库连接
    print("4. 数据库连接")
    with MockDatabase("sqlite:///test.db") as db:
        db.execute("SELECT * FROM users")
        db.execute("INSERT INTO logs VALUES ('test')")
        print(f"  执行了 {len(db.operations)} 条 SQL")
    print(f"  连接状态: {db.connected}")
    print()

    # 5. suppress
    print("5. suppress 静默异常")
    with suppress(ZeroDivisionError):
        result = 1 / 0
        print("这行不会执行")
    print("  suppress 成功吞掉了 ZeroDivisionError")


async def demo_async():
    print("\n=== 异步上下文管理器演示 ===\n")

    async with managed_http_session() as session:
        for i in range(3):
            await asyncio.sleep(0.01)
            session["requests"] += 1
            print(f"  发送请求 #{session['requests']}")


def main():
    demo_sync()
    asyncio.run(demo_async())


if __name__ == "__main__":
    main()

---

六、踩坑实录 1：contextmanager 里忘记 try/finally

# 错误：如果 with 块里抛异常，finally 不会执行，资源泄漏！
@contextmanager
def bad_db():
    conn = connect()
    yield conn
    conn.close()  # 异常时这行不会执行！

# 正确：必须用 try/finally
@contextmanager
def good_db():
    conn = connect()
    try:
        yield conn
    finally:
        conn.close()  # 无论如何都会执行

---

七、踩坑实录 2：exit 返回值的含义

class SilentContext:
    def __enter__(self): return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type == ValueError:
            print("吞掉 ValueError")
            return True   # 返回 True = 异常被处理，不再传播
        return False      # 返回 False = 异常继续传播

with SilentContext():
    raise ValueError("被吞掉了")  # 不会往外抛

with SilentContext():
    raise TypeError("会往外抛")   # 这个会传播

---

八、踩坑实录 3：嵌套 with 语句的简化写法

# 冗长写法
with open("a.txt") as f1:
    with open("b.txt") as f2:
        with open("c.txt") as f3:
            ...

# Python 3.10+ 的简化写法
with (
    open("a.txt") as f1,
    open("b.txt") as f2,
    open("c.txt") as f3,
):
    ...

# 或者
with open("a.txt") as f1, open("b.txt") as f2:
    ...

---

总结

上下文管理器的核心：

1. 实现 __enter__/__exit__ 或用 @contextmanager + yield
2. yield 前是获取，yield 后（finally）是释放
3. contextlib 提供了 suppress、redirect_stdout、asynccontextmanager 等开箱即用工具
4. 异步版本用 @asynccontextmanager 和 async with
5. 任何"获取-使用-释放"模式，首选上下文管理器