Python 内置模块 contextlib

基础概念

with 语句回顾，在了解 contextlib 之前，先回顾一下 with 语句的基本用法：

# 传统的文件操作
file = open('test.txt', 'w')
file.write('Hello World')
file.close()

# 使用 with 语句
with open('test.txt', 'w') as file:
    file.write('Hello World')
# 文件会自动关闭

contextmanager 装饰器

基本用法

contextmanager 装饰器可以将生成器函数转换为上下文管理器：

from contextlib import contextmanager

@contextmanager
def my_context():
    print("进入上下文")
    yield "返回值"
    print("退出上下文")

with my_context() as value:
    print(f"在上下文中使用: {value}")

异常处理

contextmanager 可以处理上下文中的异常：

from contextlib import contextmanager

@contextmanager
def handle_errors():
    try:
        print("开始执行")
        yield
    except ValueError as e:
        print(f"捕获到异常: {e}")
    finally:
        print("清理资源")

with handle_errors():
    raise ValueError("测试异常")

数据库连接管理示例

from contextlib import contextmanager
import sqlite3

@contextmanager
def database_connection(db_name):
    conn = sqlite3.connect(db_name)
    try:
        print(f"连接到数据库: {db_name}")
        yield conn
    finally:
        conn.close()
        print("数据库连接已关闭")

with database_connection(':memory:') as conn:
    cursor = conn.cursor()
    cursor.execute('CREATE TABLE test (name TEXT)')
    cursor.execute("INSERT INTO test VALUES ('Alice')")

closing

closing 用于确保对象的 close() 方法被调用：

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://httpbin.org/json')) as response:
    data = response.read()
    print(f"响应长度: {len(data)}")
# response.close() 会自动调用

suppress

suppress 用于抑制指定的异常：

from contextlib import suppress
import os

# 尝试删除文件，如果文件不存在也不会报错
with suppress(FileNotFoundError):
    os.remove('nonexistent_file.txt')
    print("文件删除成功")

print("程序继续执行")

# 抑制多种异常类型
with suppress(ValueError, TypeError):
    result = int("not_a_number")

redirect_stdout 和 redirect_stderr

这两个函数用于重定向标准输出和标准错误：

from contextlib import redirect_stdout, redirect_stderr
import sys
from io import StringIO

# 重定向标准输出
output_buffer = StringIO()
with redirect_stdout(output_buffer):
    print("这会被重定向")
    print("这也会被重定向")

captured_output = output_buffer.getvalue()
print(f"捕获的输出: {captured_output}")

# 重定向标准错误
error_buffer = StringIO()
with redirect_stderr(error_buffer):
    sys.stderr.write("错误信息\n")

captured_error = error_buffer.getvalue()
print(f"捕获的错误: {captured_error}")

ExitStack

ExitStack 允许动态管理多个上下文管理器：

from contextlib import ExitStack

# 管理多个文件
filenames = ['file1.txt', 'file2.txt', 'file3.txt']

with ExitStack() as stack:
    files = [stack.enter_context(open(fname, 'w')) for fname in filenames]
    
    for i, file in enumerate(files):
        file.write(f"内容写入文件 {i+1}\n")
# 所有文件都会自动关闭

条件性上下文管理：

from contextlib import ExitStack

def process_with_optional_logging(data, enable_logging=False):
    with ExitStack() as stack:
        if enable_logging:
            log_file = stack.enter_context(open('process.log', 'w'))
            log_file.write("开始处理数据\n")
        
        # 处理数据
        result = [x * 2 for x in data]
        
        if enable_logging:
            log_file.write(f"处理完成，结果: {result}\n")
        
        return result

# 不启用日志
result1 = process_with_optional_logging([1, 2, 3])
print(f"结果1: {result1}")

# 启用日志
result2 = process_with_optional_logging([4, 5, 6], enable_logging=True)
print(f"结果2: {result2}")

nullcontext

nullcontext 提供一个空的上下文管理器，用于条件性上下文管理：

from contextlib import nullcontext

def process_data(data, use_context=True):
    context = open('output.txt', 'w') if use_context else nullcontext()
    
    with context as file:
        for item in data:
            if file:  # 如果有文件对象
                file.write(f"{item}\n")
            else:  # 如果是 nullcontext
                print(item)

# 写入文件
process_data(['apple', 'banana', 'cherry'], use_context=True)

# 打印到控制台
process_data(['dog', 'cat', 'bird'], use_context=False)

自定义上下文管理器类

除了使用 contextlib，还可以通过实现 __enter__ 和 __exit__ 方法创建上下文管理器：

class TimerContext:
    def __init__(self, name):
        self.name = name
    
    def __enter__(self):
        import time
        self.start_time = time.time()
        print(f"开始计时: {self.name}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        end_time = time.time()
        duration = end_time - self.start_time
        print(f"结束计时: {self.name}, 耗时: {duration:.4f}秒")
        return False  # 不抑制异常

with TimerContext("数据处理"):
    sum_result = sum(range(1000000))
    print(f"计算结果: {sum_result}")

总结

contextlib 模块提供了多种创建和使用上下文管理器的方式：

• @contextmanager 装饰器：将生成器函数转换为上下文管理器
• closing：确保对象的 close() 方法被调用
• suppress：抑制指定的异常
• redirect_stdout/redirect_stderr：重定向标准输出和标准错误
• ExitStack：动态管理多个上下文管理器
• nullcontext：提供空的上下文管理器