Python @contextmanager使用教程

什么是上下文管理器？

上下文管理器是Python中用于管理资源的对象，确保资源在使用后被正确释放（如文件关闭、锁释放等）。通常使用with语句来工作。

传统方法需要实现__enter__和__exit__两个魔法方法，而@contextmanager装饰器可以简化这一过程。

@contextmanager的作用

@contextmanager是contextlib模块提供的装饰器，它允许使用生成器函数来创建上下文管理器，无需定义完整的类。

使用它可以：

简化上下文管理器的创建
减少样板代码
更清晰地管理资源
轻松处理异常

基本使用方法

下面是使用@contextmanager的基本结构：

from contextlib import contextmanager

@contextmanager
def my_context_manager():
    # 初始化资源（相当于__enter__）
    resource = acquire_resource()
    try:
        yield resource  # 资源在此处可用
    finally:
        # 清理资源（相当于__exit__）
        release_resource(resource)

实际示例：文件操作

下面是一个实用的文件操作上下文管理器：

from contextlib import contextmanager

@contextmanager
def open_file(file_path, mode='r'):
    print(f"打开文件: {file_path}")
    file = open(file_path, mode)
    try:
        yield file  # 文件对象在此处可用
    finally:
        print(f"关闭文件: {file_path}")
        file.close()

# 使用上下文管理器
with open_file('example.txt', 'w') as f:
    f.write("Hello, @contextmanager!")
    print("文件写入完成")
    
# 输出：
# 打开文件: example.txt
# 文件写入完成
# 关闭文件: example.txt

异常处理

@contextmanager能正确处理异常：

@contextmanager
def handle_exceptions():
    print("开始上下文")
    try:
        yield
    except Exception as e:
        print(f"捕获到异常: {type(e).__name__}: {e}")
    finally:
        print("清理资源")

# 测试异常处理
with handle_exceptions():
    print("在上下文中执行操作")
    raise ValueError("测试异常")

# 输出：
# 开始上下文
# 在上下文中执行操作
# 捕获到异常: ValueError: 测试异常
# 清理资源

数据库连接管理

创建数据库连接的上下文管理器：

import sqlite3
from contextlib import contextmanager

@contextmanager
def database_connection(db_path):
    conn = sqlite3.connect(db_path)
    print(f"连接数据库: {db_path}")
    try:
        yield conn
    finally:
        conn.close()
        print(f"关闭数据库连接")

# 使用数据库上下文管理器
with database_connection('test.db') as conn:
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
    cursor.execute("INSERT INTO users (name) VALUES ('Alice')")
    conn.commit()
    print("数据插入成功")

临时目录管理

创建临时工作目录的上下文管理器：

import os
import tempfile
from contextlib import contextmanager

@contextmanager
def temporary_directory():
    temp_dir = tempfile.mkdtemp()
    print(f"创建临时目录: {temp_dir}")
    try:
        yield temp_dir
    finally:
        print(f"删除临时目录: {temp_dir}")
        # 在实际应用中，这里应该添加删除目录的代码
        # import shutil
        # shutil.rmtree(temp_dir)

# 使用临时目录
with temporary_directory() as temp_dir:
    file_path = os.path.join(temp_dir, 'temp_file.txt')
    with open(file_path, 'w') as f:
        f.write("临时文件内容")
    print(f"在临时目录中创建了文件: {file_path}")

性能分析上下文管理器

创建代码执行时间分析的上下文管理器：

import time
from contextlib import contextmanager

@contextmanager
def timer_context(name="代码块"):
    start_time = time.perf_counter()
    try:
        yield
    finally:
        end_time = time.perf_counter()
        print(f"{name}执行时间: {end_time - start_time:.4f}秒")

# 测试代码执行时间
with timer_context("复杂计算"):
    total = 0
    for i in range(1000000):
        total += i
    print(f"计算结果: {total}")

# 输出：
# 计算结果: 499999500000
# 复杂计算执行时间: 0.0453秒

最佳实践

总是使用try...finally确保资源被释放
在yield语句之前放置资源获取代码
在finally块中放置资源清理代码
为上下文管理器使用描述性名称
处理可能发生的异常
考虑资源获取失败的情况

传统方式 vs @contextmanager

传统方式	@contextmanager
需要定义类	使用生成器函数
实现__enter__和__exit__方法	使用yield分隔资源获取和释放
代码量较多	代码更简洁
适合复杂场景	适合简单到中等复杂度的场景