Python 上下文管理器详解

上下文管理器的基本概念

上下文管理器是 Python 中用于管理资源的对象，它定义了进入和退出上下文时应该执行的操作。最常见的使用方式是 with 语句。

为什么需要上下文管理器

python
# 不好的做法 - 手动管理资源
file = open('example.txt', 'r')
try:
    content = file.read()
    process(content)
finally:
    file.close()

# 好的做法 - 使用上下文管理器
with open('example.txt', 'r') as file:
    content = file.read()
    process(content)
# 文件自动关闭

上下文管理器协议

上下文管理器需要实现两个方法：

__enter__(self): 进入上下文时调用
__exit__(self, exc_type, exc_val, exc_tb): 退出上下文时调用

基本实现

python
class MyContextManager:
    def __init__(self, resource):
        self.resource = resource
    
    def __enter__(self):
        print("进入上下文")
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出上下文")
        # 返回 True 表示抑制异常，False 表示传播异常
        return False

# 使用上下文管理器
with MyContextManager("资源") as resource:
    print(f"使用资源: {resource}")

# 输出:
# 进入上下文
# 使用资源: 资源
# 退出上下文

异常处理

python
class SafeContextManager:
    def __enter__(self):
        print("进入安全上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出安全上下文")
        if exc_type is not None:
            print(f"捕获到异常: {exc_type.__name__}: {exc_val}")
            # 返回 True 抑制异常
            return True
        return False

# 测试异常处理
with SafeContextManager():
    print("执行操作")
    raise ValueError("测试异常")

print("程序继续执行")

# 输出:
# 进入安全上下文
# 执行操作
# 退出安全上下文
# 捕获到异常: ValueError: 测试异常
# 程序继续执行

contextlib 模块

@contextmanager 装饰器

@contextmanager 装饰器简化了上下文管理器的创建。

python
from contextlib import contextmanager

@contextmanager
def simple_context():
    print("进入上下文")
    try:
        yield "资源"
    finally:
        print("退出上下文")

# 使用
with simple_context() as resource:
    print(f"使用资源: {resource}")

带异常处理的上下文管理器

python
from contextlib import contextmanager

@contextmanager
def error_handling_context():
    print("进入错误处理上下文")
    try:
        yield
    except ValueError as e:
        print(f"处理 ValueError: {e}")
        raise  # 重新抛出异常
    finally:
        print("清理资源")

# 使用
try:
    with error_handling_context():
        print("执行操作")
        raise ValueError("测试异常")
except ValueError:
    print("捕获到重新抛出的异常")

closing 函数

closing 函数为没有上下文管理器协议的对象创建上下文管理器。

python
from contextlib import closing

class Resource:
    def __init__(self, name):
        self.name = name
    
    def close(self):
        print(f"关闭资源: {self.name}")

# 使用 closing
with closing(Resource("数据库连接")) as resource:
    print(f"使用资源: {resource.name}")
# 资源自动关闭

suppress 函数

suppress 函数用于忽略指定的异常。

python
from contextlib import suppress

# 忽略 FileNotFoundError
with suppress(FileNotFoundError):
    with open('nonexistent.txt', 'r') as f:
        content = f.read()

print("程序继续执行")

# 忽略多个异常
with suppress(FileNotFoundError, PermissionError):
    with open('protected.txt', 'r') as f:
        content = f.read()

redirect_stdout 和 redirect_stderr

python
from contextlib import redirect_stdout, redirect_stderr
import io

# 重定向标准输出
output = io.StringIO()
with redirect_stdout(output):
    print("这条消息被重定向")

print(f"捕获的输出: {output.getvalue()}")

# 重定向标准错误
error_output = io.StringIO()
with redirect_stderr(error_output):
    print("错误消息", file=sys.stderr)

print(f"捕获的错误: {error_output.getvalue()}")

实际应用场景

1. 文件操作

python
# 自动关闭文件
with open('input.txt', 'r') as input_file:
    with open('output.txt', 'w') as output_file:
        for line in input_file:
            output_file.write(line.upper())

# 两个文件都会自动关闭

2. 数据库连接

python
import sqlite3
from contextlib import contextmanager

@contextmanager
def database_connection(db_path):
    conn = sqlite3.connect(db_path)
    try:
        yield conn
    finally:
        conn.close()

# 使用
with database_connection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users")
    results = cursor.fetchall()
# 连接自动关闭

3. 锁管理

python
import threading
from contextlib import contextmanager

class LockManager:
    def __init__(self):
        self.lock = threading.Lock()
    
    @contextmanager
    def acquire(self):
        self.lock.acquire()
        try:
            yield
        finally:
            self.lock.release()

# 使用
lock_manager = LockManager()

with lock_manager.acquire():
    # 临界区代码
    print("执行临界区操作")
# 锁自动释放

4. 临时目录

python
import tempfile
import os

# 创建临时目录
with tempfile.TemporaryDirectory() as temp_dir:
    print(f"临时目录: {temp_dir}")
    temp_file = os.path.join(temp_dir, 'temp.txt')
    with open(temp_file, 'w') as f:
        f.write("临时数据")
    # 临时目录及其内容在退出时自动删除
# 临时目录已被删除

5. 计时器

python
import time
from contextlib import contextmanager

@contextmanager
def timer(name):
    start_time = time.time()
    yield
    end_time = time.time()
    print(f"{name} 耗时: {end_time - start_time:.4f} 秒")

# 使用
with timer("数据处理"):
    data = [i ** 2 for i in range(1000000)]
    sum(data)

# 输出: 数据处理 耗时: 0.1234 秒

6. 临时改变配置

python
from contextlib import contextmanager

class Config:
    def __init__(self):
        self.debug = False
        self.verbose = False

config = Config()

@contextmanager
def temporary_config(config_obj, **kwargs):
    """临时修改配置"""
    original_values = {}
    
    # 保存原始值
    for key, value in kwargs.items():
        original_values[key] = getattr(config_obj, key)
        setattr(config_obj, key, value)
    
    try:
        yield
    finally:
        # 恢复原始值
        for key, value in original_values.items():
            setattr(config_obj, key, value)

# 使用
print(f"调试模式: {config.debug}")  # False

with temporary_config(config, debug=True, verbose=True):
    print(f"调试模式: {config.debug}")  # True
    print(f"详细模式: {config.verbose}")  # True

print(f"调试模式: {config.debug}")  # False

7. 事务管理

python
from contextlib import contextmanager

class Database:
    def __init__(self):
        self.in_transaction = False
        self.data = {}
    
    def begin_transaction(self):
        self.in_transaction = True
        self.transaction_data = self.data.copy()
    
    def commit(self):
        self.in_transaction = False
        print("事务提交")
    
    def rollback(self):
        self.in_transaction = False
        self.data = self.transaction_data
        print("事务回滚")

@contextmanager
def transaction(db):
    db.begin_transaction()
    try:
        yield db
        db.commit()
    except Exception as e:
        db.rollback()
        raise

# 使用
db = Database()

try:
    with transaction(db):
        db.data['key1'] = 'value1'
        db.data['key2'] = 'value2'
        # raise Exception("测试异常")  # 会触发回滚
except Exception as e:
    print(f"事务失败: {e}")

print(db.data)

嵌套上下文管理器

使用多个 with 语句

python
# 嵌套使用
with open('file1.txt', 'r') as f1:
    with open('file2.txt', 'r') as f2:
        content1 = f1.read()
        content2 = f2.read()
        # 处理两个文件

使用 contextlib.ExitStack

ExitStack 允许动态管理多个上下文管理器。

python
from contextlib import ExitStack

files = ['file1.txt', 'file2.txt', 'file3.txt']

with ExitStack() as stack:
    file_handles = [stack.enter_context(open(f, 'r')) for f in files]
    contents = [f.read() for f in file_handles]
    # 所有文件在退出时自动关闭

条件性上下文管理器

python
from contextlib import ExitStack, nullcontext

def get_context(use_context):
    if use_context:
        return MyContextManager("资源")
    else:
        return nullcontext()

# 条件性使用上下文管理器
with get_context(True) as resource:
    if resource is not None:
        print(f"使用资源: {resource}")
    else:
        print("不使用上下文管理器")

异步上下文管理器

异步上下文管理器协议

异步上下文管理器实现：

__aenter__(self): 异步进入上下文
__aexit__(self, exc_type, exc_val, exc_tb): 异步退出上下文

python
class AsyncContextManager:
    def __init__(self, resource):
        self.resource = resource
    
    async def __aenter__(self):
        print("异步进入上下文")
        await self.connect()
        return self.resource
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        print("异步退出上下文")
        await self.disconnect()
    
    async def connect(self):
        print("连接资源...")
    
    async def disconnect(self):
        print("断开连接...")

# 使用异步上下文管理器
async def use_async_context():
    async with AsyncContextManager("异步资源") as resource:
        print(f"使用资源: {resource}")

# 运行
import asyncio
asyncio.run(use_async_context())

asynccontextmanager 装饰器

python
from contextlib import asynccontextmanager

@asynccontextmanager
async def async_resource():
    print("获取异步资源")
    try:
        yield "异步资源"
    finally:
        print("释放异步资源")

async def use_async_resource():
    async with async_resource() as resource:
        print(f"使用资源: {resource}")

asyncio.run(use_async_resource())

最佳实践

1. 确保资源清理

python
# 好的做法 - 使用 finally 确保清理
class ResourceManager:
    def __enter__(self):
        self.resource = acquire_resource()
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        release_resource(self.resource)
        return False

2. 正确处理异常

python
# 好的做法 - 区分异常类型
class SafeContextManager:
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            print("正常退出")
        elif issubclass(exc_type, (ValueError, TypeError)):
            print(f"处理预期异常: {exc_val}")
            return True  # 抑制预期异常
        else:
            print(f"未处理异常: {exc_val}")
            return False  # 传播未处理异常

3. 提供有用的错误信息

python
class DatabaseContextManager:
    def __enter__(self):
        try:
            self.connection = connect_to_database()
            return self.connection
        except ConnectionError as e:
            raise RuntimeError(f"无法连接到数据库: {e}")
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.connection:
            try:
                self.connection.close()
            except Exception as e:
                print(f"关闭连接时出错: {e}")

4. 支持上下文管理器链

python
class ChainedContextManager:
    def __init__(self, *managers):
        self.managers = managers
    
    def __enter__(self):
        self.entered_managers = []
        try:
            for manager in self.managers:
                self.entered_managers.append(manager.__enter__())
            return self.entered_managers[-1]
        except:
            # 如果失败，清理已进入的上下文
            self.__exit__(None, None, None)
            raise
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 反向退出上下文
        for manager in reversed(self.entered_managers):
            manager.__exit__(exc_type, exc_val, exc_tb)
        return False

总结

Python 上下文管理器的核心概念：

基本协议：实现 __enter__ 和 __exit__ 方法
with 语句：自动管理资源的进入和退出
异常处理：在 __exit__ 中处理异常
contextlib 模块：简化上下文管理器的创建
实际应用：文件操作、数据库连接、锁管理等
嵌套管理：使用多个 with 语句或 ExitStack
异步支持：异步上下文管理器用于异步代码

上下文管理器的优势：

自动资源管理，避免资源泄漏
代码更简洁、更易读
异常安全，确保资源清理
支持嵌套和链式使用

掌握上下文管理器，能够编写出更安全、更优雅的 Python 代码。

Python 中的上下文管理器（Context Manager）是什么？如何使用？