1. contextlib #

本节介绍 contextlib 模块的基本定位。contextlib 是 Python 标准库中用于简化上下文管理器创建的实用工具模块。它提供了装饰器、函数和类来帮助开发者更优雅地管理资源（如文件、数据库连接、网络连接等），确保资源在使用完毕后被正确释放，避免资源泄漏。

2. 核心功能 #

本节介绍 contextlib 模块中最常用的几个核心工具。这些工具可以大大简化上下文管理器的创建和使用，让代码更加简洁和 Pythonic。

2.1. @contextmanager 装饰器 #

@contextmanager 是 contextlib 中最核心的工具，它允许你使用生成器函数来创建上下文管理器，而不需要编写完整的类。这种方式比传统的类实现更加简洁，特别适合简单的资源管理场景。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义一个模拟的资源获取和释放函数（实际项目中替换为真实逻辑）
def acquire_resource(*args, **kwargs):
    # 模拟获取资源
    print(f"获取资源，参数: {args}, {kwargs}")
    return f"资源_{args[0] if args else 'default'}"

# 定义一个模拟的资源释放函数
def release_resource(resource):
    # 模拟释放资源
    print(f"释放资源: {resource}")

# 使用 @contextmanager 装饰器创建上下文管理器
@contextmanager
def managed_resource(*args, **kwargs):
    # 设置代码（相当于 __enter__ 方法）
    resource = acquire_resource(*args, **kwargs)
    try:
        # 将资源提供给 with 块使用
        yield resource
    finally:
        # 清理代码（相当于 __exit__ 方法，确保总是执行）
        release_resource(resource)

# 定义主函数演示用法
def main():
    # 使用示例：在 with 块中使用资源
    with managed_resource("test", mode="read") as r:
        # 在这里使用资源 r
        print(f"使用资源: {r}")
        print("资源使用中...")

# 脚本入口
if __name__ == '__main__':
    main()

示例：文件操作上下文管理器

这个示例展示如何创建一个自定义的文件操作上下文管理器，虽然 Python 内置的 open() 已经支持 with 语句，但这里演示了自定义实现的原理。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 使用 @contextmanager 创建文件操作的上下文管理器
@contextmanager
def open_file(path, mode):
    # 打开文件
    f = open(path, mode)
    try:
        # 将文件对象提供给 with 块
        yield f
    finally:
        # 确保文件被关闭
        f.close()

# 定义主函数
def main():
    # 使用自定义的文件上下文管理器
    with open_file('test.txt', 'w') as f:
        # 写入内容
        f.write('Hello, World!')

    # 读取刚才写入的内容
    with open_file('test.txt', 'r') as f:
        # 读取内容
        content = f.read()
        print(f"读取的内容: {content}")

# 脚本入口
if __name__ == '__main__':
    main()

2.2. closing() - 自动关闭对象 #

closing() 函数为具有 close() 方法但不实现上下文管理器协议的对象创建上下文管理器。这在处理第三方库的对象时特别有用，确保资源被正确关闭。

# 导入 closing 函数
from contextlib import closing
# 导入 urllib.request 中的 urlopen
from urllib.request import urlopen

# 定义主函数
def main():
    # 使用 closing() 为 urlopen 创建上下文管理器
    with closing(urlopen('http://www.python.org')) as page:
        # 读取页面内容
        content = page.read()
        # 页面会在 with 块结束时自动关闭
        print(f"页面内容长度: {len(content)} 字节")

    # 注意：现代 Python 中，urlopen 已经支持 with 语句
    # 所以上面的代码等价于：
    with urlopen('http://www.python.org') as page:
        content = page.read()
        print(f"直接使用 with 语句，内容长度: {len(content)} 字节")

# 脚本入口
if __name__ == '__main__':
    main()

2.3. suppress() - 抑制特定异常 #

suppress() 函数用于临时抑制指定的异常，这在需要忽略特定错误但不想使用 try/except 块的情况下很有用。它让代码更加简洁，专注于主要逻辑。

# 导入 suppress 函数
from contextlib import suppress
# 导入 os 模块
import os

# 定义主函数
def main():
    # 删除文件，如果文件不存在也不报错
    with suppress(FileNotFoundError):
        # 尝试删除可能不存在的文件
        os.remove('somefile.tmp')
        print("文件删除成功")

    # 上面的代码等价于：
    try:
        os.remove('somefile.tmp')
        print("文件删除成功")
    except FileNotFoundError:
        # 忽略文件不存在的错误
        pass

    print("程序继续执行")

# 脚本入口
if __name__ == '__main__':
    main()

2.4. nullcontext() - 空上下文管理器 #

nullcontext() 提供一个什么都不做的上下文管理器，用于代码需要上下文管理器但实际不需要的特殊情况。这在条件性使用上下文管理器时特别有用。

# 导入 nullcontext 函数
from contextlib import nullcontext

# 定义一个需要上下文管理器的数据处理函数
def process_data(data, context_manager=None):
    # 如果没有提供上下文管理器，使用空上下文管理器
    if context_manager is None:
        context_manager = nullcontext()

    # 在上下文中处理数据
    with context_manager:
        # 处理数据
        return data.upper()

# 定义主函数
def main():
    # 使用空上下文管理器
    result = process_data('hello')
    print(f"处理结果: {result}")

    # 也可以传入真实的上下文管理器
    from contextlib import contextmanager

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

    # 使用真实的上下文管理器
    result = process_data('world', log_context())
    print(f"处理结果: {result}")

# 脚本入口
if __name__ == '__main__':
    main()

2.5. ExitStack() - 动态管理多个上下文 #

ExitStack() 是一个强大的工具，用于管理多个上下文管理器的动态栈。它特别适用于需要动态数量上下文管理器的情况，或者需要条件性地添加上下文管理器的场景。

# 导入 ExitStack 类
from contextlib import ExitStack

# 定义一个处理多个文件的函数
def process_files(filenames):
    # 使用 ExitStack 管理多个文件
    with ExitStack() as stack:
        # 动态打开多个文件
        files = [stack.enter_context(open(fname)) for fname in filenames]
        # 所有文件都会在 with 块结束时自动关闭
        # 处理文件内容
        contents = [f.read() for f in files]
        return contents

# 定义主函数
def main():
    # 创建一些测试文件
    test_files = ['file1.txt', 'file2.txt', 'file3.txt']

    # 创建测试文件
    for filename in test_files:
        with open(filename, 'w') as f:
            f.write(f"这是 {filename} 的内容")

    try:
        # 处理多个文件
        contents = process_files(test_files)
        for i, content in enumerate(contents):
            print(f"文件 {i+1} 内容: {content}")
    finally:
        # 清理测试文件
        import os
        for filename in test_files:
            if os.path.exists(filename):
                os.remove(filename)

# 脚本入口
if __name__ == '__main__':
    main()

更复杂的 ExitStack 示例：

这个示例展示如何更灵活地使用 ExitStack，包括条件性地添加上下文管理器。

# 导入 ExitStack 类
from contextlib import ExitStack

# 定义一个模拟的文件打开函数
def open_file(fname):
    # 模拟文件打开
    print(f"打开文件: {fname}")
    return f"文件对象_{fname}"

# 定义一个模拟的数据库连接函数
def get_db_connection():
    # 模拟数据库连接
    print("建立数据库连接")
    return "数据库连接对象"

# 定义一个模拟的条件变量
condition = True

# 定义主函数
def main():
    # 使用 ExitStack 管理多个上下文管理器
    with ExitStack() as stack:
        # 动态添加多个文件上下文管理器
        resources = [
            stack.enter_context(open_file(fname))
            for fname in ['a.txt', 'b.txt', 'c.txt']
        ]

        # 还可以条件性地添加上下文管理器
        if condition:
            # 根据条件添加数据库连接
            db_connection = stack.enter_context(get_db_connection())
            print(f"使用数据库连接: {db_connection}")

        # 使用资源
        print(f"管理的资源数量: {len(resources)}")
        print("所有资源都会自动管理")

        # 模拟一些操作
        for resource in resources:
            print(f"使用资源: {resource}")

# 脚本入口
if __name__ == '__main__':
    main()

3. 高级用法 #

本节介绍 contextlib 的一些高级用法，包括嵌套上下文管理器、带参数的上下文管理器以及错误处理和资源清理等技巧。这些用法可以帮助你创建更复杂和灵活的上下文管理器。

3.1. 嵌套上下文管理器 #

嵌套上下文管理器允许你组合多个上下文管理器，创建更复杂的资源管理逻辑。这种方式特别适合需要多个相关资源的场景。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义一个模拟的数据库对象
class MockDatabase:
    def transaction(self):
        # 返回一个模拟的事务上下文管理器
        return self

    def __enter__(self):
        print("开始数据库事务")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            print("提交数据库事务")
        else:
            print("回滚数据库事务")

# 创建一个模拟的数据库实例
db = MockDatabase()

# 定义数据库事务的上下文管理器
@contextmanager
def database_transaction(db):
    # 在数据库事务中执行操作
    with db.transaction():
        yield

# 定义日志操作的上下文管理器
@contextmanager
def log_operation(operation_name):
    # 记录操作开始
    print(f"开始操作: {operation_name}")
    try:
        yield
    finally:
        # 记录操作完成
        print(f"完成操作: {operation_name}")

# 定义主函数
def main():
    # 组合使用多个上下文管理器
    with log_operation("数据处理"), database_transaction(db):
        # 执行数据库操作
        print("执行数据处理...")
        # 模拟一些数据处理逻辑
        result = "处理完成"
        print(f"结果: {result}")

# 脚本入口
if __name__ == '__main__':
    main()

3.2. 带参数的上下文管理器 #

带参数的上下文管理器可以根据传入的参数调整其行为，这使得上下文管理器更加灵活和可重用。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义一个带参数的计时器上下文管理器
@contextmanager
def timer(name):
    # 导入 time 模块
    import time
    # 记录开始时间
    start = time.time()
    try:
        # 执行被计时的操作
        yield
    finally:
        # 计算并打印耗时
        end = time.time()
        print(f"{name} 耗时: {end - start:.2f} 秒")

# 定义一个模拟的耗时计算函数
def heavy_computation():
    # 模拟耗时操作
    import time
    time.sleep(0.1)  # 休眠 0.1 秒
    return "计算完成"

# 定义主函数
def main():
    # 使用带参数的计时器
    with timer("计算任务"):
        # 执行耗时操作
        result = heavy_computation()
        print(f"计算结果: {result}")

    # 使用不同的名称再次计时
    with timer("另一个任务"):
        # 执行另一个操作
        print("执行另一个任务...")
        time.sleep(0.05)  # 休眠 0.05 秒

# 脚本入口
if __name__ == '__main__':
    main()

3.3. 错误处理和资源清理 #

错误处理和资源清理是上下文管理器的核心功能。良好的错误处理可以确保即使在出现异常的情况下，资源也能被正确清理。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义一个安全的操作上下文管理器
@contextmanager
def safe_operation():
    try:
        # 执行操作
        yield
    except Exception as e:
        # 捕获并处理异常
        print(f"操作失败: {e}")
        # 可以选择重新抛出异常或处理它
        raise
    finally:
        # 确保资源清理总是执行
        print("资源清理完成")

# 定义一个模拟的风险操作函数
def risky_operation():
    # 模拟可能失败的操作
    import random
    if random.random() < 0.5:
        raise ValueError("随机错误")
    print("操作成功")

# 定义主函数
def main():
    try:
        # 使用安全的操作上下文管理器
        with safe_operation():
            risky_operation()
    except Exception as e:
        print(f"捕获到异常: {e}")

    print("程序继续执行")

# 脚本入口
if __name__ == '__main__':
    main()

4. 实际应用场景 #

本节介绍 contextlib 在实际项目中的常见应用场景。这些示例展示了如何使用上下文管理器来解决真实的编程问题，包括数据库事务管理、临时目录处理、状态管理等。

4.1. 数据库事务管理 #

数据库事务管理是上下文管理器的典型应用场景。通过使用上下文管理器，可以确保数据库事务在成功时被提交，在失败时被回滚，并且连接总是被正确关闭。

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义一个模拟的数据库会话类
class MockSession:
    def __init__(self):
        self.committed = False
        self.rolled_back = False
        self.closed = False

    def add(self, obj):
        print(f"添加对象: {obj}")

    def commit(self):
        print("提交事务")
        self.committed = True

    def rollback(self):
        print("回滚事务")
        self.rolled_back = True

    def close(self):
        print("关闭会话")
        self.closed = True

# 定义一个模拟的用户类
class User:
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return f"User({self.name})"

# 定义数据库事务的上下文管理器
@contextmanager
def db_transaction(session):
    try:
        # 提供会话给 with 块使用
        yield session
        # 如果没有异常，提交事务
        session.commit()
    except Exception:
        # 如果有异常，回滚事务
        session.rollback()
        # 重新抛出异常
        raise
    finally:
        # 确保会话被关闭
        session.close()

# 定义主函数
def main():
    # 创建模拟的数据库会话
    session = MockSession()

    try:
        # 使用数据库事务上下文管理器
        with db_transaction(session) as s:
            # 在事务中添加用户
            s.add(User(name='John'))
            print("用户添加成功")

        # 检查事务状态
        print(f"事务已提交: {session.committed}")
        print(f"事务已回滚: {session.rolled_back}")
        print(f"会话已关闭: {session.closed}")

    except Exception as e:
        print(f"操作失败: {e}")

# 脚本入口
if __name__ == '__main__':
    main()

4.2. 临时目录处理 #

临时目录处理是另一个常见的应用场景。通过使用上下文管理器，可以确保临时目录在使用完毕后被正确清理，避免磁盘空间浪费。

# 导入必要的模块
import tempfile
import shutil
import os

# 导入 contextmanager 装饰器
from contextlib import contextmanager

# 定义临时目录的上下文管理器
@contextmanager
def temp_directory():
    # 创建临时目录
    dirpath = tempfile.mkdtemp()
    try:
        # 提供临时目录路径给 with 块使用
        yield dirpath
    finally:
        # 确保临时目录被删除
        shutil.rmtree(dirpath)

# 定义主函数
def main():
    # 使用临时目录上下文管理器
    with temp_directory() as temp_dir:
        # 在临时目录中工作
        print(f"使用临时目录: {temp_dir}")

        # 创建测试文件
        test_file = os.path.join(temp_dir, "test.txt")
        with open(test_file, "w") as f:
            f.write("test content")

        # 读取文件内容
        with open(test_file, "r") as f:
            content = f.read()
            print(f"文件内容: {content}")

        # 检查临时目录是否存在
        print(f"临时目录存在: {os.path.exists(temp_dir)}")

    # 检查临时目录是否已被删除
    print(f"临时目录已被删除: {not os.path.exists(temp_dir)}")

# 脚本入口
if __name__ == '__main__':
    main()

4.3. 状态管理 #

状态管理是上下文管理器的另一个重要应用。通过使用上下文管理器，可以临时修改系统状态（如环境变量、配置等），并在使用完毕后恢复原始状态。

# 导入 contextmanager 装饰器
from contextlib import contextmanager
# 导入 os 模块
import os

# 定义环境变量管理的上下文管理器
@contextmanager
def set_temp_env(var_name, value):
    # 保存原始值
    original = os.environ.get(var_name)
    # 设置新的环境变量值
    os.environ[var_name] = value
    try:
        # 在修改后的环境中执行操作
        yield
    finally:
        # 恢复原始值
        if original is None:
            # 如果原来没有这个环境变量，删除它
            os.environ.pop(var_name, None)
        else:
            # 如果原来有这个环境变量，恢复它的值
            os.environ[var_name] = original

# 定义一个模拟的调试模式运行函数
def run_debug_mode():
    # 检查调试环境变量
    debug_value = os.environ.get('DEBUG', '0')
    print(f"当前 DEBUG 值: {debug_value}")

    if debug_value == '1':
        print("运行在调试模式")
    else:
        print("运行在正常模式")

# 定义主函数
def main():
    # 显示原始环境变量值
    print(f"原始 DEBUG 值: {os.environ.get('DEBUG', '未设置')}")

    # 使用环境变量管理上下文管理器
    with set_temp_env('DEBUG', '1'):
        # 在这个块中，DEBUG 环境变量为 '1'
        print("在修改后的环境中:")
        run_debug_mode()

    # 检查环境变量是否已恢复
    print(f"恢复后的 DEBUG 值: {os.environ.get('DEBUG', '未设置')}")

    # 再次运行，应该使用原始值
    print("在原始环境中:")
    run_debug_mode()

# 脚本入口
if __name__ == '__main__':
    main()

5. 最佳实践 #

本节总结使用 contextlib 和创建上下文管理器时的一些最佳实践。遵循这些实践可以帮助你创建更可靠、更易维护的代码。

1. 总是使用 try/finally：确保在 @contextmanager 中使用 try/finally 来保证清理代码总是执行

2. 正确处理异常：考虑是否应该在上下文管理器中捕获和处理异常

3. 保持简洁：上下文管理器应该专注于单一资源的管程

4. 提供有用的错误信息：当资源获取失败时，提供清晰的错误信息

6. 总结 #

contextlib 模块提供了强大而灵活的工具来创建和使用上下文管理器，使得资源管理变得更加简洁和可靠。通过使用这些工具，你可以：

• 减少样板代码
• 提高代码的可读性
• 确保资源正确释放
• 实现更优雅的错误处理

上下文管理器是 Python 中资源管理的最佳实践，而 contextlib 让创建和使用上下文管理器变得更加简单和优雅。