Python 上下文管理器的高级应用与自定义实现

Python 的 with 语句是处理资源管理的黄金标准，最常见的应用场景就是文件操作。当我们使用 with open() 时，无论代码块中是否发生异常，文件都会被正确关闭。这种自动化的资源管理大大提升了代码的可靠性。

上下文管理器本质上是一个实现了 __enter__() 和 __exit__() 方法的对象。当进入 with 语句时，Python 会调用 __enter__() 方法，该方法的返回值会被绑定到 as 后的变量上。当代码块执行完毕或发生异常时，__exit__() 会被调用，负责清理工作。

让我们从一个实用的例子开始：创建一个计时器上下文管理器。这在性能分析和调试时非常有用。

import time

class Timer:
    def __init__(self, name='Operation'):
        self.name = name
        self.start_time = None
        self.elapsed = None

    def __enter__(self):
        self.start_time = time.time()
        print(f'[{self.name}] 开始执行')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.time() - self.start_time
        print(f'[{self.name}] 执行完成，耗时: {self.elapsed:.3f} 秒')
        if exc_type is not None:
            print(f'[{self.name}] 发生异常: {exc_type.__name__}')
        return False

# 使用示例
with Timer('数据处理'):
    time.sleep(1)
    # 模拟耗时操作
    result = sum(range(1000000))

这个计时器会在进入和退出上下文时自动打印时间信息。__exit__() 方法的返回值很重要：返回 True 表示异常已被处理，返回 False 则异常会继续传播。

接下来，让我们看一个更复杂的应用：数据库连接的上下文管理器。在实际项目中，手动管理数据库连接很容易导致连接泄漏，使用上下文管理器可以完美解决这个问题。

import sqlite3

class DatabaseConnection:
    def __init__(self, db_path):
        self.db_path = db_path
        self.connection = None
        self.cursor = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_path)
        self.cursor = self.connection.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            # 没有异常，提交事务
            self.connection.commit()
        else:
            # 发生异常，回滚事务
            self.connection.rollback()
            print(f'数据库操作失败，事务已回滚: {exc_val}')

        self.cursor.close()
        self.connection.close()
        return False

# 使用示例
with DatabaseConnection('example.db') as cursor:
    cursor.execute('CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)')
    cursor.execute('INSERT INTO users (name) VALUES (?)', ('Alice',))
        # 如果这里抛出异常，事务会自动回滚

这个数据库上下文管理器展示了上下文管理器的强大之处：自动提交或回滚事务，确保连接始终被关闭。这样即使代码中发生异常，数据库状态也能保持一致。

Python 还提供了一个更简洁的方式创建上下文管理器：contextlib 模块。特别是 @contextmanager 装饰器，它允许我们用生成器函数来定义上下文管理器，无需创建类。

from contextlib import contextmanager
import tempfile
import os

@contextmanager
def temporary_file(content=None):
    '创建临时文件的上下文管理器'
    fd, path = tempfile.mkstemp()
    try:
        if content:
            with os.fdopen(fd, 'w') as f:
                f.write(content)
        else:
            os.close(fd)
            fd = None

        yield path  # 返回文件路径给 with 块使用

    finally:
        # 确保文件描述符和文件都被清理
        if fd is not None:
            os.close(fd)
        if os.path.exists(path):
            os.unlink(path)

# 使用示例
with temporary_file('临时内容') as temp_path:
    print(f'临时文件路径: {temp_path}')
    with open(temp_path, 'r') as f:
        print(f'文件内容: {f.read()}')
# 文件在这里已自动删除

这个临时文件管理器确保文件在使用后立即被删除，非常适合非常适合处理敏感数据或需要临时存储的场景。yield。语句之前的代码相当于 __enter__()，之后的代码相当于 __exit__()。

让我们再看一个更高级的应用：重试机制的上下文管理器。在网络请求或可能失败的操作中，自动重试是一个常见需求。

from contextlib import contextmanager
import time
import random

@contextmanager
def retry(max_attempts=3, delay=1.0):
    '自动重试的上下文管理器'
    attempts = 0
    last_exception = None

    while attempts < max_attempts:
        attempts += 1
        try:
            yield
            return  # 成功执行，直接返回

        except Exception as e:
            last_exception = e
            if attempts < max_attempts:
                print(f'尝试 {attempts}/{max_attempts} 失败，{delay} 秒后重试...')
                time.sleep(delay)
            else:
                print(f'所有 {max_attempts} 次尝试都失败了')

    # 所有尝试都失败，重新抛出最后的异常
    raise last_exception

# 使用示例
@retry(max_attempts=3, delay=1.0)
def unreliable_operation():
    '模拟可能失败的操作'
    if random.random() < 0.7:
        raise ValueError('随机失败')
    print('操作成功！')
    return 'success'

unreliable_operation()

这个重试管理器展示了上下文管理器的灵活性。通过装饰器和上下文管理器的组合，我们创建了一个既优雅又实用的解决方案。

contextlib 模块还提供了几个现成的上下文管理器工具。closing() 可以用于任何有 close() 方法的对象；redirect_stdout() 可以临时重定向标准输出；suppress() 可以抑制特定异常。

from contextlib import closing, redirect_stdout, suppress
import io

# closing() 示例
with closing(open('example.txt', 'w')) as f:
    f.write('内容')
# 文件自动关闭

# redirect_stdout() 示例
buffer = io.StringIO()
with redirect_stdout(buffer):
    print('这不会打印到控制台')
print(f'捕获的输出: {buffer.getvalue()}')

# suppress() 示例
with suppress(FileNotFoundError):
    os.remove('不存在的文件')
# 不会抛出异常

上下文管理器还可以嵌套使用，Python 的 contextlib.nested() 或直接用 with 语句的多重语法可以实现这一点。

# 嵌套上下文管理器示例
with Timer('总操作'), DatabaseConnection('example.db') as cursor:
    cursor.execute('SELECT * FROM users')
    results = cursor.fetchall()
    print(f'查询到 {len(results)} 条记录')

在实际开发中，合理使用上下文管理器可以显著提升代码质量和可维护性。它不仅仅是一种语法糖，更是一种设计哲学：将资源管理的责任封装在专门的对象中，让业务逻辑更加清晰。

总结来说，Python 的上下文管理器是处理资源管理的强大工具。无论是自定义类还是使用 @contextmanager 装饰器，都能帮助我们创建更优雅、更安全的代码。建议在以下场景优先考虑使用上下文管理器：文件操作、数据库连接、网络连接、锁管理、临时资源、计时和性能分析。

掌握上下文管理器是成为高级 Python 开发者的重要一步。希望本文的示例和解释能够帮助你更好地理解和使用这个强大的特性。