当前位置：首页 > news >正文

Python @contextmanager 装饰器

news 2025/9/13 8:17:10

Python @contextmanager 装饰器：优雅的资源管理之道

@contextmanager 是 Python 中一个极其强大的装饰器，它允许你使用生成器函数来创建上下文管理器，让资源管理变得异常简洁和优雅。本文将带你深入理解这个装饰器的原理和用法。

一、什么是上下文管理器？

在了解 @contextmanager 之前，我们先回顾一下上下文管理器的概念。上下文管理器用于管理资源（如文件、锁、数据库连接等）的获取和释放，确保资源被正确清理。

传统方式：使用类实现

class FileManager:def __init__(self, filename, mode):self.filename = filenameself.mode = modeself.file = Nonedef __enter__(self):self.file = open(self.filename, self.mode)return self.filedef __exit__(self, exc_type, exc_val, exc_tb):if self.file:self.file.close()# 使用
with FileManager('test.txt', 'w') as f:f.write('Hello World')

这种方式需要实现 __enter__ 和 __exit__ 两个方法，代码量较多。

二、@contextmanager 的基本用法

@contextmanager 装饰器来自 contextlib 模块，它让我们可以用更简洁的方式创建上下文管理器。

基本语法

from contextlib import contextmanager@contextmanager
def simple_context():#  setup 代码 (相当于 __enter__)print("进入上下文")yield "资源对象"#  teardown 代码 (相当于 __exit__)print("退出上下文")# 使用
with simple_context() as resource:print(f"使用资源: {resource}")print("在上下文中执行操作")

输出：

进入上下文
使用资源: 资源对象
在上下文中执行操作
退出上下文

三、实际应用案例

案例1：文件操作

from contextlib import contextmanager@contextmanager
def open_file(filename, mode):"""自动关闭文件的上下文管理器"""try:print(f"打开文件: {filename}")file = open(filename, mode)yield file  # 将文件对象返回给 with 语句finally:print(f"关闭文件: {filename}")file.close()# 使用
with open_file('example.txt', 'w') as f:f.write('Hello, World!')print("文件写入完成")# 输出:
# 打开文件: example.txt
# 文件写入完成
# 关闭文件: example.txt

案例2：数据库连接管理

from contextlib import contextmanager
import sqlite3@contextmanager
def database_connection(db_path):"""自动管理数据库连接的上下文管理器"""conn = Nonetry:conn = sqlite3.connect(db_path)print("数据库连接已建立")yield connprint("事务提交成功")conn.commit()except Exception as e:if conn:print("发生错误，回滚事务")conn.rollback()raise efinally:if conn:conn.close()print("数据库连接已关闭")# 使用
try: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')")print("数据库操作完成")
except Exception as e:print(f"错误: {e}")

案例3：计时器

from contextlib import contextmanager
import time@contextmanager
def timer(name="操作"):"""测量代码执行时间的上下文管理器"""start_time = time.time()try:yieldfinally:end_time = time.time()print(f"{name}耗时: {end_time - start_time:.4f} 秒")# 使用
with timer("数据计算"):# 模拟耗时操作result = sum(i * i for i in range(1000000))print(f"计算结果: {result}")

案例4：临时目录管理

from contextlib import contextmanager
import tempfile
import os@contextmanager
def temporary_directory():"""创建临时目录并在退出时自动清理"""temp_dir = tempfile.mkdtemp()print(f"创建临时目录: {temp_dir}")try:yield temp_dirfinally:# 清理临时目录import shutilshutil.rmtree(temp_dir)print(f"已删除临时目录: {temp_dir}")# 使用
with temporary_directory() as temp_dir:# 在临时目录中工作temp_file = os.path.join(temp_dir, 'temp.txt')with open(temp_file, 'w') as f:f.write('临时文件内容')print(f"在目录 {temp_dir} 中创建了文件")

四、错误处理

@contextmanager 可以很好地处理异常：

from contextlib import contextmanager@contextmanager
def error_handling_context():"""演示错误处理的上下文管理器"""try:print("资源初始化")yield "资源"except ValueError as e:print(f"捕获到 ValueError: {e}")# 可以重新抛出或其他处理raiseexcept Exception as e:print(f"捕获到其他异常: {e}")raisefinally:print("资源清理")# 测试不同类型的异常
print("=== 测试 ValueError ===")
try:with error_handling_context() as res:print(f"使用 {res}")raise ValueError("这是一个值错误")
except:passprint("\n=== 测试 RuntimeError ===")
try:with error_handling_context() as res:print(f"使用 {res}")raise RuntimeError("这是一个运行时错误")
except:passprint("\n=== 测试正常流程 ===")
with error_handling_context() as res:print(f"使用 {res}")

五、高级用法

1. 嵌套上下文管理器

from contextlib import contextmanager@contextmanager
def indent_context(level):"""缩进上下文"""print("  " * level + "↘ 进入")try:yieldfinally:print("  " * level + "↗ 退出")# 嵌套使用
with indent_context(0):print("第一层")with indent_context(1):print("第二层")with indent_context(2):print("第三层")

2. 带参数的上下文管理器

from contextlib import contextmanager@contextmanager
def managed_resource(resource_name, **kwargs):"""通用的资源管理器"""print(f"初始化资源: {resource_name}")resource = f"{resource_name}_instance"  # 这里应该是实际的资源初始化try:yield resourceexcept Exception as e:print(f"资源 {resource_name} 使用过程中出错: {e}")raisefinally:print(f"清理资源: {resource_name}")# 使用不同类型的资源
with managed_resource("数据库连接") as db:print(f"使用 {db} 执行查询")with managed_resource("文件句柄", mode='w') as file:print(f"向 {file} 写入数据")

3. 组合多个上下文管理器

from contextlib import contextmanager@contextmanager
def database_transaction(db_config):"""数据库事务管理器"""print("开始事务")try:yield "数据库连接"print("提交事务")except:print("回滚事务")raisefinally:print("关闭连接")@contextmanager 
def log_operation(operation_name):"""操作日志管理器"""print(f"开始操作: {operation_name}")try:yieldprint(f"操作成功: {operation_name}")except:print(f"操作失败: {operation_name}")raise# 组合使用
def complex_operation():with log_operation("用户注册"):with database_transaction({"host": "localhost"}) as db:print(f"使用 {db} 执行用户注册操作")# 模拟业务逻辑return "注册成功"# 执行
result = complex_operation()
print(f"结果: {result}")

六、与传统方式的对比

传统类方式 vs @contextmanager

# 方式1：使用类
class Timer:def __init__(self, name):self.name = namedef __enter__(self):self.start = time.time()return selfdef __exit__(self, *args):self.end = time.time()print(f"{self.name} 耗时: {self.end - self.start:.4f}s")# 方式2：使用 @contextmanager
@contextmanager
def timer(name):start = time.time()yieldend = time.time()print(f"{name} 耗时: {end - start:.4f}s")# 使用对比
with Timer("类方式") as t:time.sleep(0.1)with timer("装饰器方式"):time.sleep(0.1)

七、最佳实践

1. 总是使用 try-finally

from contextlib import contextmanager@contextmanager
def good_example():resource = acquire_resource()  # 获取资源try:yield resourcefinally:release_resource(resource)  # 确保资源释放@contextmanager
def bad_example():resource = acquire_resource()yield resource  # 如果这里发生异常，资源不会被释放！release_resource(resource)

2. 适当的错误处理

from contextlib import contextmanager@contextmanager
def robust_context():resource = Nonetry:resource = initialize_expensive_resource()yield resourceexcept CustomException as e:handle_special_case(e)except Exception as e:log_error(e)raise  # 重新抛出未知异常finally:if resource is not None:cleanup_resource(resource)

3. 提供有用的上下文信息

from contextlib import contextmanager@contextmanager
def informative_context(operation_name):"""提供详细信息的上下文管理器"""print(f"🚀 开始: {operation_name}")start_time = time.time()try:yieldstatus = "成功"except Exception as e:status = f"失败: {e}"raisefinally:end_time = time.time()duration = end_time - start_timeprint(f"✅ {operation_name} {status} (耗时: {duration:.2f}s)")