上下文管理器 (Context Manager)与 with 语句

with 语句是 Python 中一种优雅的资源管理方式，它可以自动分配和释放资源，确保即使在发生异常时也能正确清理资源。最常见的用法是文件操作，但它的应用远不止于此。

1. 为什么要用 `with`？

在不使用 with 的情况下，我们需要手动关闭资源。如果操作过程中发生异常，可能导致资源未能正确关闭。

传统写法（不推荐）

f = open('test.txt', 'w')
try:
    f.write('Hello, World!')
except Exception as e:
    print(f"An error occurred: {e}")
finally:
    f.close() # 必须在 finally 中关闭，防止异常导致文件未关闭

这种写法略显繁琐，容易遗漏 finally 块。

`with` 写法（推荐）

with open('test.txt', 'w') as f:
    f.write('Hello, World!')
# 出代码块后，文件会自动关闭，即使发生异常也不例外

代码更加简洁、清晰，且更安全。

2. 上下文管理器协议

任何实现了 __enter__ 和 __exit__ 方法的对象都可以作为上下文管理器使用。

__enter__(self): 进入 with 代码块前调用，返回值会赋值给 as 后的变量。
__exit__(self, exc_type, exc_value, traceback): 离开 with 代码块时调用，用于清理资源。

自定义上下文管理器

我们可以创建一个计时器，统计代码块的执行时间。

import time

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

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.time()
        self.interval = self.end - self.start
        print(f"代码执行耗时: {self.interval:.4f} 秒")
        # 如果返回 True，则会抑制异常，不向上抛出
        # return False 

# 使用自定义上下文管理器
with Timer() as t:
    for i in range(1000000):
        pass

3. `contextlib` 模块

Python 提供了 contextlib 模块，方便我们更轻易地实现上下文管理器，特别是 contextmanager 装饰器。

使用 `@contextmanager`

我们可以将一个生成器函数快速转化为上下文管理器。

from contextlib import contextmanager

@contextmanager
def my_open(filename, mode):
    print("准备打开文件...")
    f = open(filename, mode)
    try:
        yield f # yield 之前是 __enter__ 的逻辑
    finally:
        print("正在关闭文件...")
        f.close() # yield 之后是 __exit__ 的逻辑

# 使用
with my_open('test.txt', 'w') as f:
    f.write('Hello via contextlib!')

这种写法利用了生成器的特性，yield 之前的代码在进入 with 时执行，yield 之后的代码在离开 with 时执行。

4. 常见应用场景

数据库连接

自动提交事务或回滚，并在结束后关闭连接。

线程锁 (Lock)

import threading

lock = threading.Lock()

# 自动获取锁和释放锁
with lock:
    # 临界区代码
    pass

临时修改环境

例如临时改变日志级别或当前工作目录。

import os
from contextlib import contextmanager

@contextmanager
def change_dir(destination):
    try:
        cwd = os.getcwd()
        os.chdir(destination)
        yield
    finally:
        os.chdir(cwd)

with change_dir('/tmp'):
    print(os.getcwd()) # /tmp
print(os.getcwd()) # 恢复原目录

总结

with 语句用于自动管理资源（如文件、锁、网络连接）。
它通过上下文管理器协议（__enter__ 和 __exit__）工作。
使用 contextlib.contextmanager 可以简化上下文管理器的编写。
始终优先使用 with 语句来处理需要关闭或清理的资源。

下一步

如果说上下文管理器是处理代码块的“前后逻辑”，那么还有一种更灵活的方式，可以在函数调用的前后织入代码，那就是装饰器。

👉 前往下一章：装饰器 (Decorators)