上下文管理器

概述

上下文管理器（Context Manager）是 Python 中用于管理资源的机制，确保资源在使用后被正确释放。通过 with 语句，上下文管理器可以在进入和退出代码块时自动执行设置和清理操作，即使代码块中发生异常也能保证清理逻辑被执行。

基础概念

with 语句的工作原理

with 语句在进入代码块时调用 __enter__，退出时调用 __exit__：

with expression as variable:
    # 代码块
    pass

# 等价于
manager = expression
variable = manager.__enter__()
try:
    # 代码块
    pass
finally:
    manager.__exit__(exc_type, exc_val, exc_tb)

enter 和 exit

• __enter__：进入 with 块时调用，返回值赋给 as 后的变量
• __exit__：退出 with 块时调用，接收异常信息（如果没有异常，三个参数都是 None）
• __exit__ 返回 True 表示异常已被处理，不再向外传播

快速上手

类实现上下文管理器

class DatabaseConnection:
    """数据库连接上下文管理器"""
    def __init__(self, url):
        self.url = url

    def __enter__(self):
        self.conn = connect(self.url)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.close()
        return False  # 不抑制异常

with DatabaseConnection('postgresql://...') as conn:
    conn.execute('SELECT 1')

contextmanager 装饰器

使用 @contextmanager 装饰器可以用生成器更简洁地实现上下文管理器：

from contextlib import contextmanager
import time

@contextmanager
def timer(name):
    """计时上下文管理器"""
    start = time.perf_counter()
    yield  # yield 之前的代码相当于 __enter__，之后相当于 __exit__
    elapsed = time.perf_counter() - start
    print(f"{name}: {elapsed:.3f}s")

with timer("query"):
    db.query("SELECT * FROM users")

详细用法

contextlib 工具集

suppress

忽略指定异常：

from contextlib import suppress

# 忽略 FileNotFoundError
with suppress(FileNotFoundError):
    os.remove('temp.txt')
# 等价于
try:
    os.remove('temp.txt')
except FileNotFoundError:
    pass

closing

自动调用 close() 方法：

from contextlib import closing

with closing(urllib.request.urlopen(url)) as response:
    data = response.read()
# 退出时自动调用 response.close()

redirect_stdout / redirect_stderr

重定向标准输出：

from contextlib import redirect_stdout
import io

# 将 print 输出捕获到字符串
output = io.StringIO()
with redirect_stdout(output):
    print("这行不会显示在控制台")
    print("而是被捕获到 output 中")

result = output.getvalue()  # 获取捕获的内容

ExitStack

动态管理多个上下文：

from contextlib import ExitStack

# 动态打开多个文件
filenames = ['a.txt', 'b.txt', 'c.txt']

with ExitStack() as stack:
    files = [stack.enter_context(open(f)) for f in filenames]
    # 所有文件会在退出时自动关闭
    for f in files:
        print(f.read())

异常处理

@contextmanager
def transaction(db):
    """数据库事务上下文管理器"""
    try:
        db.begin()
        yield db
        db.commit()      # 正常退出时提交
    except Exception:
        db.rollback()    # 异常时回滚
        raise            # 重新抛出异常

with transaction(db) as t:
    t.execute("INSERT INTO users VALUES (1, 'Alice')")
    t.execute("INSERT INTO orders VALUES (100, 1)")

嵌套上下文

# 同时打开输入和输出文件
with open('input.txt') as fin, open('output.txt', 'w') as fout:
    for line in fin:
        fout.write(line.upper())

# 等价于
with open('input.txt') as fin:
    with open('output.txt', 'w') as fout:
        for line in fin:
            fout.write(line.upper())

异步上下文管理器

class AsyncDB:
    """异步数据库连接"""
    async def __aenter__(self):
        self.conn = await connect()
        return self.conn

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.conn.close()

async with AsyncDB() as conn:
    await conn.execute('SELECT 1')

使用 @asynccontextmanager：

from contextlib import asynccontextmanager

@asynccontextmanager
async def async_timer(name):
    """异步计时器"""
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"{name}: {elapsed:.3f}s")

async with async_timer("async query"):
    await db.query("SELECT * FROM users")

常见场景

场景一：文件操作

@contextmanager
def safe_write(filepath):
    """安全写入文件：先写入临时文件，成功后再重命名"""
    tmp_path = filepath + '.tmp'
    try:
        with open(tmp_path, 'w') as f:
            yield f
        os.rename(tmp_path, filepath)  # 原子操作
    except Exception:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)  # 失败时删除临时文件
        raise

with safe_write('config.json') as f:
    json.dump(config, f)
# 只有写入成功才会替换原文件

场景二：临时修改环境

import os

@contextmanager
def temp_env(**kwargs):
    """临时修改环境变量"""
    old_values = {}
    for key, value in kwargs.items():
        old_values[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for key, old_value in old_values.items():
            if old_value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old_value

with temp_env(DATABASE_URL="postgresql://test", DEBUG="1"):
    # 在这个代码块中环境变量被临时修改
    run_tests()
# 退出后环境变量恢复原值

场景三：性能分析

import cProfile
import pstats

@contextmanager
def profile(output_file=None):
    """性能分析上下文管理器"""
    profiler = cProfile.Profile()
    profiler.enable()
    try:
        yield profiler
    finally:
        profiler.disable()
        stats = pstats.Stats(profiler)
        stats.sort_stats('cumulative')
        if output_file:
            stats.dump_stats(output_file)
        else:
            stats.print_stats(20)  # 打印前 20 行

with profile("profile.stats"):
    expensive_operation()

场景四：锁管理

import threading

@contextmanager
def acquire_lock(lock, timeout=5):
    """带超时的锁管理"""
    acquired = lock.acquire(timeout=timeout)
    if not acquired:
        raise TimeoutError("获取锁超时")
    try:
        yield
    finally:
        lock.release()

lock = threading.Lock()
with acquire_lock(lock, timeout=10):
    # 在锁保护下执行操作
    update_shared_resource()

注意事项

• __exit__ 返回 True 会吞掉异常，除非有明确需求，否则应返回 False
• 使用 @contextmanager 时，yield 之前的代码相当于 __enter__，之后的代码相当于 __exit__
• @contextmanager 中 yield 之后的清理代码应放在 try/finally 中，确保即使 yield 本身抛出异常也能清理
• 异步上下文管理器使用 async with，不能用普通的 with
• 不要在 __exit__ 中执行可能抛出异常的操作，这会掩盖原始异常

进阶用法

可重入的上下文管理器

class ReentrantManager:
    """可重入的上下文管理器"""
    _depth = 0

    def __enter__(self):
        if self._depth == 0:
            print("首次进入")
        self._depth += 1
        return self

    def __exit__(self, *exc):
        self._depth -= 1
        if self._depth == 0:
            print("最终退出")
        return False

上下文管理器工厂

class DBTransaction:
    """可配置的事务管理器"""
    def __init__(self, db, isolation_level='READ COMMITTED'):
        self.db = db
        self.isolation_level = isolation_level

    def __enter__(self):
        self.db.execute(f"SET TRANSACTION ISOLATION LEVEL {self.isolation_level}")
        self.db.begin()
        return self.db

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.db.commit()
        else:
            self.db.rollback()
        return False

# 使用
with DBTransaction(db, 'SERIALIZABLE') as t:
    t.execute("UPDATE accounts SET balance = balance - 100")
    t.execute("UPDATE accounts SET balance = balance + 100")

使用 contextlib.nullcontext

Python 3.7+ 提供了空上下文管理器，用于条件性使用上下文：

from contextlib import nullcontext

def process(lock=None):
    """可选的锁保护"""
    ctx = lock if lock else nullcontext()
    with ctx:
        # 有锁时加锁执行，无锁时直接执行
        do_work()

上下文管理器与装饰器结合

from functools import wraps

def contextmanager_to_decorator(cm_factory):
    """将上下文管理器工厂转换为装饰器"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with cm_factory():
                return func(*args, **kwargs)
        return wrapper
    return decorator

# 使用
@contextmanager_to_decorator(timer)
def my_function():
    # 函数执行时自动计时
    expensive_computation()