Python 上下文管理器：自定义上下文管理器实现

在 Python 开发中，我们经常需要管理资源，比如打开文件、操作数据库连接、处理网络请求等。这些场景有一个共同特点：使用前需要获取资源，使用后必须释放资源。如果忘记释放，或者释放过程中遇到异常，就会导致资源泄漏。

上下文管理器就是为解决这一问题而设计的机制。它通过 with 语句确保资源在任何情况下都能被正确清理，让代码更简洁、更安全。

一、上下文管理器的基本概念

Python 中的 with 语句背后的机制就是上下文管理器协议。任何实现了 __enter__ 和 __exit__ 方法的对象都可以作为上下文管理器使用。

当你执行以下代码时：

with open("example.txt", "r") as f:
    content = f.read()

Python 内部会做这些事情：调用 f.__enter__() 获取返回值（这里就是 f 本身），将返回值赋值给 f；代码块执行完毕后，自动调用 f.__exit__() **关闭文件`。

这种模式的核心价值在于：无论代码块是否发生异常， __exit__ 都会被执行，从而保证资源被正确释放。

二、自定义上下文管理器：类实现方式

最直接的实现方式是创建一个类，在类中定义 __enter__ 和 __exit__ 方法。

2.1 基础模板

以下是一个标准的上下文管理器类结构：

class MyContextManager:
    def __enter__(self):
        # 进入上下文时执行
        # 返回值会赋值给 with 语句中的 as 变量
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # 退出上下文时执行
        # exc_type/value/traceback 记录了是否发生异常
        # 返回 True 表示异常已被处理，不向外传播
        # 返回 False 或 None 表示异常未处理，会继续向外抛出
        pass

2.2 实际示例：计时器

假设你需要测量代码块的执行时间，可以这样实现：

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end = time.perf_counter()
        self.elapsed = self.end - self.start
        print(f"执行耗时: {self.elapsed:.4f} 秒")

# 使用方式
with Timer() as t:
    # 这里放置要计时的代码
    total = sum(range(1000000))

# 可以访问计时结果
print(f"本次耗时: {t.elapsed:.4f} 秒")

执行这段代码，你会看到类似这样的输出：

执行耗时: 0.0231 秒
本次耗时: 0.0231 秒

2.3 处理异常

__exit__ 方法的三个参数可以帮你检测和处理异常：

class SafeContext:
    def __enter__(self):
        print("资源已获取")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print("正常退出，无异常")
        else:
            print(f"检测到异常: {exc_type.__name__}: {exc_value}")
            # 返回 True 吞掉异常
            return True

with SafeContext() as ctx:
    print("代码块执行中...")
    raise ValueError("出错了！")

print("程序继续执行")

输出如下：

资源已获取
代码块执行中...
检测到异常: ValueError: 出错了！
程序继续执行

注意：由于 __exit__ 返回了 ValueError，异常被捕获处理了，所以程序能够继续执行。

三、使用 @contextmanager 实现更简洁

对于简单的场景，每次都写一个类显得有些繁琐。Python 在 contextlib 模块中提供了 @contextmanager 装饰器，可以用生成器函数的方式创建上下文管理器。

3.1 基本用法

from contextlib import contextmanager

@contextmanager
def my_context():
    # 这部分代码相当于 __enter__ 的逻辑
    print("进入上下文")
    try:
        yield  # 暂停在这里，将控制权交给 with 语句块
    finally:
        # 这部分代码相当于 __exit__ 的逻辑
        print("退出上下文")

使用方式完全一样：

with my_context() as value:
    print("代码块执行中")

输出如下：

进入上下文
代码块执行中
退出上下文

3.2 实际示例：临时目录

下面实现一个自动创建和清理临时目录的上下文管理器：

import os
import tempfile
from contextlib import contextmanager

@contextmanager
def temp_directory():
    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    print(f"创建临时目录: {temp_dir}")

    try:
        yield temp_dir  # 将临时目录路径yield给调用者
    finally:
        # 清理临时目录
        for filename in os.listdir(temp_dir):
            os.remove(os.path.join(temp_dir, filename))
        os.rmdir(temp_dir)
        print(f"已清理临时目录: {temp_dir}")

# 使用方式
with temp_directory() as td:
    # 在临时目录中创建文件
    test_file = os.path.join(td, "test.txt")
    with open(test_file, "w") as f:
        f.write("测试数据")
    print(f"在 {td} 创建了测试文件")

print("临时目录已被自动清理")

3.3 异常处理在 @contextmanager 中

使用 @contextmanager 时，异常处理逻辑需要放在 try-finally 结构中：

@contextmanager
def transaction(db_connection):
    db_connection.begin()
    try:
        yield
    except Exception as e:
        db_connection.rollback()
        print(f"事务回滚: {e}")
    else:
        db_connection.commit()
        print("事务提交成功")

四、常见应用场景

4.1 数据库事务管理

@contextmanager
def db_transaction(conn):
    conn.begin()
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()

# 使用
with db_transaction(database) as db:
    db.execute("INSERT INTO users VALUES (...)")
    db.execute("UPDATE accounts SET balance = ...")

4.2 锁管理

import threading

@contextmanager
def thread_lock(lock):
    lock.acquire()
    try:
        yield
    finally:
        lock.release()

# 使用
lock = threading.Lock()
with thread_lock(lock):
    # 访问共享资源的代码
    shared_data += 1

4.3 临时修改环境设置

@contextmanager
def temp_env(**kwargs):
    old_env = {}
    for key, value in kwargs.items():
        old_env[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for key, value in old_env.items():
            if value is None:
                del os.environ[key]
            else:
                os.environ[key] = value

# 使用：临时修改时区设置
with temp_env(TZ='Asia/Shanghai'):
    print(datetime.now())

4.4 性能追踪

@contextmanager
def track_performance(operation_name):
    import time
    start = time.time()
    metrics = {}
    yield metrics
    elapsed = time.time() - start
    metrics["operation"] = operation_name
    metrics["duration_ms"] = round(elapsed * 1000, 2)
    print(f"[{operation_name}] 耗时: {metrics['duration_ms']}ms")

# 使用
with track_performance("数据处理") as m:
    result = process_large_dataset()
print(f"metrics: {m}")

五、进阶技巧

5.1 多重上下文嵌套

Python 支持在同一行使用多个 with 语句，这在处理需要同时管理多个资源的场景时非常方便：

with open("input.txt", "r") as fin, open("output.txt", "w") as fout:
    content = fin.read()
    fout.write(content.upper())

5.2 使用 contextlib.ExitStack

当上下文管理器数量在运行时才能确定时，ExitStack 是最佳选择：

from contextlib import ExitStack

files = []
try:
    with ExitStack() as stack:
        # 根据条件打开多个文件
        for path in file_paths:
            f = stack.enter_context(open(path, "r"))
            files.append(f)
        # 处理文件...
        # ExitStack 会自动清理所有打开的文件
finally:
    # 此时所有文件已经被关闭
    pass

5.3 让现有对象支持上下文管理

对于已经存在的类，可以通过 contextlib.ContextDecorator 添加上下文管理能力：

from contextlib import ContextDecorator

class Timer(ContextDecorator):
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, *args):
        self.elapsed = time.perf_counter() - self.start
        print(f"耗时: {self.elapsed:.4f}s")

# 作为装饰器使用
@Timer()
def long_running_task():
    time.sleep(1)

# 作为上下文管理器使用
with Timer():
    time.sleep(0.5)

六、最佳实践总结

选择实现方式的建议：

编写上下文管理器时遵循的原则：

1. 始终使用 try-finally：确保退出逻辑总是执行，即使发生异常
2. 正确处理异常：决定是吞掉异常还是让它继续传播
3. 清理顺序与获取顺序相反：这符合资源管理的最佳实践
4. 避免在 yield 之后放置代码：@contextmanager 中 yield 之后的代码只有在异常时才会执行

上下文管理器是 Python 中优雅的资源管理机制。掌握它不仅能让代码更简洁，还能从根本上避免资源泄漏问题。建议在实际开发中多思考哪些场景适合使用上下文管理器，逐步培养这种编码习惯。