深入解析Python上下文管理器与contextlib

上下文管理器是Python中管理资源的重要机制，它确保资源在使用后被正确释放。本文将深入解析上下文管理器的工作原理和contextlib模块的高级用法。

上下文管理器基础

上下文管理器通过__enter__和__exit__方法实现，配合with语句使用，确保资源的正确获取和释放。

上下文管理器核心实现

"""
Python上下文管理器与contextlib深入解析
包含自定义上下文管理器、contextlib工具、异步上下文等
"""

import time
import threading
import logging
from typing import Optional, Type, Any
from contextlib import (
    contextmanager, closing, suppress, redirect_stdout,
    redirect_stderr, ExitStack, asynccontextmanager
)
from io import StringIO
import asyncio

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ============ 基础上下文管理器 ============
class DatabaseConnection:
    """数据库连接上下文管理器"""
    
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        self.connection = None
        self.is_connected = False
    
    def __enter__(self):
        """进入上下文"""
        logger.info(f"连接到数据库: {self.connection_string}")
        self.connection = {"status": "connected", "id": id(self)}
        self.is_connected = True
        time.sleep(0.1)  # 模拟连接时间
        logger.info("数据库连接成功")
        return self.connection
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文"""
        if self.is_connected:
            logger.info("关闭数据库连接")
            self.connection = None
            self.is_connected = False
        
        # 处理异常
        if exc_type:
            logger.error(f"发生异常: {exc_val}")
            # 返回True表示异常已处理，不再传播
            return False


class FileLock:
    """文件锁上下文管理器"""
    
    def __init__(self, filename: str):
        self.filename = filename
        self.lock = threading.Lock()
        self.acquired = False
    
    def __enter__(self):
        logger.info(f"尝试获取文件锁: {self.filename}")
        self.lock.acquire()
        self.acquired = True
        logger.info(f"文件锁已获取: {self.filename}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.acquired:
            self.lock.release()
            logger.info(f"文件锁已释放: {self.filename}")


class Timer:
    """计时器上下文管理器"""
    
    def __init__(self, name: str = "Operation"):
        self.name = name
        self.start_time = None
        self.elapsed = None
    
    def __enter__(self):
        self.start_time = time.time()
        logger.info(f"[{self.name}] 开始计时")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.time() - self.start_time
        logger.info(f"[{self.name}] 耗时: {self.elapsed:.4f}秒")


class Transaction:
    """事务上下文管理器"""
    
    def __init__(self, db_connection):
        self.db = db_connection
        self.committed = False
    
    def __enter__(self):
        logger.info("开始事务")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            logger.info("事务回滚")
            self.rollback()
        else:
            if not self.committed:
                self.commit()
    
    def commit(self):
        logger.info("提交事务")
        self.committed = True
    
    def rollback(self):
        logger.info("回滚事务")


# ============ 使用contextmanager装饰器 ============
@contextmanager
def managed_resource(resource_name: str):
    """管理资源的上下文管理器"""
    logger.info(f"获取资源: {resource_name}")
    resource = {"name": resource_name, "status": "active"}
    try:
        yield resource
    finally:
        logger.info(f"释放资源: {resource_name}")


@contextmanager
def temporary_attribute(obj, attr_name, value):
    """临时修改对象属性"""
    old_value = getattr(obj, attr_name, None)
    setattr(obj, attr_name, value)
    try:
        yield
    finally:
        if old_value is None:
            delattr(obj, attr_name)
        else:
            setattr(obj, attr_name, old_value)


@contextmanager
def suppress_exceptions(*exceptions):
    """抑制指定异常"""
    try:
        yield
    except exceptions as e:
        logger.warning(f"抑制异常: {e}")


@contextmanager
def retry_on_failure(max_attempts: int = 3, delay: float = 1.0):
    """失败重试上下文"""
    for attempt in range(1, max_attempts + 1):
        try:
            yield attempt
            break
        except Exception as e:
            logger.warning(f"第{attempt}次尝试失败: {e}")
            if attempt == max_attempts:
                raise
            time.sleep(delay)


# ============ ExitStack高级用法 ============
def demonstrate_exit_stack():
    """演示ExitStack"""
    print("\n" + "="*60)
    print("ExitStack演示")
    print("="*60)
    
    with ExitStack() as stack:
        # 动态管理多个上下文
        files = []
        for i in range(3):
            # 模拟打开文件
            ctx = managed_resource(f"resource_{i}")
            resource = stack.enter_context(ctx)
            files.append(resource)
        
        logger.info(f"管理的资源: {[f['name'] for f in files]}")
        
        # 可以动态添加回调
        stack.callback(logger.info, "ExitStack清理完成")
    
    print("所有资源已清理")


# ============ suppress用法 ============
def demonstrate_suppress():
    """演示suppress"""
    print("\n" + "="*60)
    print("suppress演示")
    print("="*60)
    
    # 抑制特定异常
    with suppress(FileNotFoundError):
        with open("nonexistent_file.txt", "r") as f:
            content = f.read()
    
    print("FileNotFoundError被抑制，程序继续执行")
    
    # 抑制多个异常
    with suppress(ZeroDivisionError, ValueError):
        result = 1 / 0
        number = int("not_a_number")
    
    print("多个异常被抑制")


# ============ redirect_stdout用法 ============
def demonstrate_redirect():
    """演示输出重定向"""
    print("\n" + "="*60)
    print("输出重定向演示")
    print("="*60)
    
    # 捕获stdout
    output = StringIO()
    with redirect_stdout(output):
        print("这条消息被重定向到StringIO")
        print("这是第二条消息")
    
    captured = output.getvalue()
    print(f"捕获的输出:\n{captured}")


# ============ closing用法 ============
class Resource:
    """需要关闭的资源"""
    def __init__(self, name):
        self.name = name
        self.closed = False
    
    def close(self):
        logger.info(f"关闭资源: {self.name}")
        self.closed = True
    
    def do_something(self):
        if self.closed:
            raise RuntimeError("资源已关闭")
        return f"使用资源: {self.name}"


def demonstrate_closing():
    """演示closing"""
    print("\n" + "="*60)
    print("closing演示")
    print("="*60)
    
    with closing(Resource("test_resource")) as resource:
        result = resource.do_something()
        logger.info(result)
    
    print(f"资源已关闭: {resource.closed}")


# ============ 异步上下文管理器 ============
class AsyncDatabaseConnection:
    """异步数据库连接"""
    
    def __init__(self, dsn: str):
        self.dsn = dsn
        self.connection = None
    
    async def __aenter__(self):
        logger.info(f"异步连接到: {self.dsn}")
        await asyncio.sleep(0.1)
        self.connection = {"dsn": self.dsn, "status": "connected"}
        return self.connection
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        logger.info("异步关闭连接")
        self.connection = None


@asynccontextmanager
async def async_managed_resource(name: str):
    """异步资源管理器"""
    logger.info(f"异步获取资源: {name}")
    resource = {"name": name}
    try:
        yield resource
    finally:
        await asyncio.sleep(0.05)
        logger.info(f"异步释放资源: {name}")


async def demonstrate_async_context():
    """演示异步上下文管理器"""
    print("\n" + "="*60)
    print("异步上下文管理器演示")
    print("="*60)
    
    async with AsyncDatabaseConnection("postgresql://localhost/db") as conn:
        logger.info(f"使用连接: {conn}")
    
    async with async_managed_resource("async_resource") as resource:
        logger.info(f"使用资源: {resource}")


# ============ 综合示例 ============
class ConfigManager:
    """配置管理器 - 展示上下文管理器的实际应用"""
    
    def __init__(self):
        self.config = {}
        self._backup = None
    
    @contextmanager
    def temporary_config(self, **kwargs):
        """临时配置上下文"""
        self._backup = self.config.copy()
        self.config.update(kwargs)
        logger.info(f"临时配置: {kwargs}")
        try:
            yield self.config
        finally:
            self.config = self._backup
            self._backup = None
            logger.info("配置已恢复")
    
    @contextmanager
    def atomic_update(self):
        """原子更新上下文"""
        backup = self.config.copy()
        try:
            yield self
            logger.info("配置原子更新成功")
        except Exception as e:
            self.config = backup
            logger.error(f"配置更新失败，已回滚: {e}")
            raise


def main():
    """主函数"""
    print("="*60)
    print("Python上下文管理器与contextlib深入解析")
    print("="*60)
    
    # 1. 基础上下文管理器
    print("\n【基础上下文管理器演示】")
    with DatabaseConnection("postgresql://localhost/mydb") as conn:
        logger.info(f"使用连接: {conn}")
    
    # 2. 文件锁
    print("\n【文件锁演示】")
    with FileLock("data.txt"):
        logger.info("执行受保护的操作")
    
    # 3. 计时器
    print("\n【计时器演示】")
    with Timer("数据处理"):
        time.sleep(0.2)
    
    # 4. 事务
    print("\n【事务演示】")
    with Transaction({"name": "test_db"}) as tx:
        logger.info("执行业务操作")
        tx.commit()
    
    # 5. contextmanager装饰器
    print("\n【contextmanager装饰器演示】")
    with managed_resource("database_pool") as resource:
        logger.info(f"使用资源: {resource}")
    
    # 6. 临时属性修改
    print("\n【临时属性修改演示】")
    class MyClass:
        value = 10
    
    obj = MyClass()
    logger.info(f"原始值: {obj.value}")
    with temporary_attribute(obj, "value", 100):
        logger.info(f"临时值: {obj.value}")
    logger.info(f"恢复后: {obj.value}")
    
    # 7. 异常抑制
    demonstrate_suppress()
    
    # 8. 输出重定向
    demonstrate_redirect()
    
    # 9. closing
    demonstrate_closing()
    
    # 10. ExitStack
    demonstrate_exit_stack()
    
    # 11. 配置管理器
    print("\n【配置管理器演示】")
    config_mgr = ConfigManager()
    config_mgr.config = {"debug": False, "timeout": 30}
    
    with config_mgr.temporary_config(debug=True, timeout=60):
        logger.info(f"临时配置: {config_mgr.config}")
    
    logger.info(f"恢复后配置: {config_mgr.config}")
    
    # 12. 异步上下文
    asyncio.run(demonstrate_async_context())
    
    print("\n" + "="*60)
    print("上下文管理器总结")
    print("="*60)
    print("1. __enter__/__exit__: 基础实现方式")
    print("2. @contextmanager: 简化实现")
    print("3. ExitStack: 动态管理多个上下文")
    print("4. suppress: 抑制特定异常")
    print("5. redirect_stdout: 重定向输出")
    print("6. closing: 确保close被调用")
    print("7. 异步上下文: async with支持")
    print("="*60)


if __name__ == "__main__":
    main()

上下文管理器执行流程

关键要点

1. 资源管理：确保资源正确获取和释放
2. 异常处理：在__exit__中处理或传播异常
3. 简化实现：使用@contextmanager装饰器
4. 组合使用：ExitStack管理多个上下文
5. 异步支持：async with和@asynccontextmanager

上下文管理器是Python资源管理的最佳实践，合理使用可以显著提升代码的健壮性和可读性。