1.5. 上下文管理器

上下文管理器是 Python 中管理资源的优雅方式，通过 with 语句确保资源被正确释放。

1.5.1. 基本概念

1.5.1.1. with 语句原理

# with 语句等价于：
# manager = context_manager()
# value = manager.__enter__()
# try:
#     ... use value ...
# finally:
#     manager.__exit__(exc_type, exc_value, traceback)

# 最常见的例子：文件操作
with open('file.txt', 'w') as f:
    f.write('Hello')
# 文件自动关闭，即使发生异常

# ⚠️ 不使用 with 的风险
f = open('file.txt')
# 如果这里发生异常...
data = f.read()
f.close()  # 可能永远不会执行！

1.5.2. 实现上下文管理器

1.5.2.1. 类方式

class Timer:
    """计时上下文管理器"""
    
    def __init__(self, name="Timer"):
        self.name = name
    
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self  # 返回值绑定到 as 后的变量
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        self.elapsed = time.perf_counter() - self.start
        print(f"{self.name}: {self.elapsed:.4f}s")
        return False  # 不抑制异常

with Timer("Processing") as t:
    import time
    time.sleep(0.5)
# Processing: 0.5012s

print(f"Elapsed: {t.elapsed:.4f}s")  # 可以在外部访问

1.5.2.2. 使用 contextlib

from contextlib import contextmanager

@contextmanager
def timer(name="Timer"):
    """生成器方式的上下文管理器"""
    import time
    start = time.perf_counter()
    try:
        yield  # yield 之前是 __enter__，之后是 __exit__
    finally:
        elapsed = time.perf_counter() - start
        print(f"{name}: {elapsed:.4f}s")

with timer("Quick operation"):
    pass  # Quick operation: 0.0000s

1.5.2.3. 异常处理

class SuppressException:
    """抑制特定异常的上下文管理器"""
    
    def __init__(self, *exceptions):
        self.exceptions = exceptions
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 返回 True 抑制异常，False 传播异常
        if exc_type is not None and issubclass(exc_type, self.exceptions):
            print(f"Suppressed: {exc_type.__name__}: {exc_val}")
            return True
        return False

with SuppressException(FileNotFoundError, PermissionError):
    open('nonexistent_file.txt')
# Suppressed: FileNotFoundError: [Errno 2] No such file or directory...

# 标准库提供了 suppress
from contextlib import suppress

with suppress(FileNotFoundError):
    open('nonexistent_file.txt')
# 静默失败

1.5.3. 实用上下文管理器

1.5.3.1. 临时修改

from contextlib import contextmanager
import os
import sys

@contextmanager
def temporary_env(**environ):
    """临时设置环境变量"""
    old_environ = dict(os.environ)
    os.environ.update(environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(old_environ)

print(os.environ.get('DEBUG'))  # None
with temporary_env(DEBUG='1', API_KEY='secret'):
    print(os.environ.get('DEBUG'))  # 1
print(os.environ.get('DEBUG'))  # None

@contextmanager
def temporary_chdir(path):
    """临时切换目录"""
    old_cwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old_cwd)

@contextmanager
def capture_stdout():
    """捕获标准输出"""
    from io import StringIO
    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        yield sys.stdout
    finally:
        sys.stdout = old_stdout

with capture_stdout() as output:
    print("This is captured")

captured = output.getvalue()
print(f"Captured: {captured!r}")  # Captured: 'This is captured\n'

1.5.3.2. 数据库事务

@contextmanager
def transaction(connection):
    """数据库事务管理"""
    cursor = connection.cursor()
    try:
        yield cursor
        connection.commit()  # 成功则提交
    except Exception:
        connection.rollback()  # 失败则回滚
        raise
    finally:
        cursor.close()

# 使用示例
# with transaction(conn) as cursor:
#     cursor.execute("INSERT INTO users VALUES (?, ?)", (1, 'Alice'))
#     cursor.execute("INSERT INTO logs VALUES (?)", ('User created',))

1.5.3.3. 锁管理

import threading
from contextlib import contextmanager

class LockManager:
    """可超时的锁管理器"""
    
    def __init__(self, lock, timeout=None):
        self.lock = lock
        self.timeout = timeout
        self.acquired = False
    
    def __enter__(self):
        self.acquired = self.lock.acquire(timeout=self.timeout or -1)
        if not self.acquired:
            raise TimeoutError("Could not acquire lock")
        return self
    
    def __exit__(self, *args):
        if self.acquired:
            self.lock.release()
        return False

lock = threading.Lock()
with LockManager(lock, timeout=5):
    # 临界区
    pass

1.5.4. 多个上下文管理器

# 嵌套写法
with open('input.txt') as f_in:
    with open('output.txt', 'w') as f_out:
        f_out.write(f_in.read())

# 简洁写法（Python 3.1+）
with open('input.txt') as f_in, open('output.txt', 'w') as f_out:
    f_out.write(f_in.read())

# 括号写法（Python 3.10+）
with (
    open('input.txt') as f_in,
    open('output.txt', 'w') as f_out,
    Timer() as t,
):
    f_out.write(f_in.read())

1.5.4.1. 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, f in enumerate(files):
        f.write(f"Content for file {i}")

# 条件性使用上下文管理器
def process_file(path, use_backup=False):
    with ExitStack() as stack:
        f = stack.enter_context(open(path))
        
        if use_backup:
            backup = stack.enter_context(open(path + '.bak', 'w'))
            backup.write(f.read())
            f.seek(0)
        
        return f.read()

1.5.5. 异步上下文管理器

import asyncio

class AsyncTimer:
    """异步计时器"""
    
    async def __aenter__(self):
        import time
        self.start = time.perf_counter()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        import time
        self.elapsed = time.perf_counter() - self.start
        print(f"Elapsed: {self.elapsed:.4f}s")
        return False

async def main():
    async with AsyncTimer():
        await asyncio.sleep(1)
    # Elapsed: 1.0012s

# asyncio.run(main())

# 使用 contextlib 的异步版本
from contextlib import asynccontextmanager

@asynccontextmanager
async def async_lock(lock, timeout=None):
    """带超时的异步锁"""
    try:
        await asyncio.wait_for(lock.acquire(), timeout)
        yield
    finally:
        lock.release()

1.5.6. 最佳实践

使用场景

资源管理：文件、网络连接、数据库连接
状态临时修改：环境变量、工作目录、全局配置
锁和同步：线程锁、分布式锁
事务管理：数据库事务、文件操作
计时和监控：性能测量、日志记录

设计原则

__enter__ 返回有用的值：通常是 self 或管理的资源
__exit__ 处理清理逻辑：即使发生异常也要执行
谨慎抑制异常：只在明确知道安全时返回 True
使用 @contextmanager 简化：简单场景更清晰

常见错误

# ❌ 忘记返回 False（或不返回）
def __exit__(self, *args):
    self.cleanup()
    return True  # 这会抑制所有异常！

# ✅ 正确做法
def __exit__(self, *args):
    self.cleanup()
    return False  # 或不返回（等同于 return None/False）