上下文管理器
Python上下文管理器详解：__enter__、__exit__、contextlib。
概述
上下文管理器（Context Manager）是 Python 中用于管理资源的机制，确保资源在使用后被正确释放。通过 with 语句，上下文管理器可以在进入和退出代码块时自动执行设置和清理操作，即使代码块中发生异常也能保证清理逻辑被执行。
基础概念
with 语句的工作原理
with 语句在进入代码块时调用 __enter__，退出时调用 __exit__：
with expression as variable:
# 代码块
pass
# 等价于
manager = expression
variable = manager.__enter__()
try:
# 代码块
pass
finally:
manager.__exit__(exc_type, exc_val, exc_tb)enter 和 exit
__enter__：进入 with 块时调用，返回值赋给 as 后的变量__exit__：退出 with 块时调用，接收异常信息（如果没有异常，三个参数都是 None）__exit__返回 True 表示异常已被处理，不再向外传播
快速上手
类实现上下文管理器
class DatabaseConnection:
"""数据库连接上下文管理器"""
def __init__(self, url):
self.url = url
def __enter__(self):
self.conn = connect(self.url)
return self.conn
def __exit__(self, exc_type, exc_val, exc_tb):
self.conn.close()
return False # 不抑制异常
with DatabaseConnection('postgresql://...') as conn:
conn.execute('SELECT 1')contextmanager 装饰器
使用 @contextmanager 装饰器可以用生成器更简洁地实现上下文管理器：
from contextlib import contextmanager
import time
@contextmanager
def timer(name):
"""计时上下文管理器"""
start = time.perf_counter()
yield # yield 之前的代码相当于 __enter__，之后相当于 __exit__
elapsed = time.perf_counter() - start
print(f"{name}: {elapsed:.3f}s")
with timer("query"):
db.query("SELECT * FROM users")详细用法
contextlib 工具集
suppress
忽略指定异常：
from contextlib import suppress
# 忽略 FileNotFoundError
with suppress(FileNotFoundError):
os.remove('temp.txt')
# 等价于
try:
os.remove('temp.txt')
except FileNotFoundError:
passclosing
自动调用 close() 方法：
from contextlib import closing
with closing(urllib.request.urlopen(url)) as response:
data = response.read()
# 退出时自动调用 response.close()redirect_stdout / redirect_stderr
重定向标准输出：
from contextlib import redirect_stdout
import io
# 将 print 输出捕获到字符串
output = io.StringIO()
with redirect_stdout(output):
print("这行不会显示在控制台")
print("而是被捕获到 output 中")
result = output.getvalue() # 获取捕获的内容ExitStack
动态管理多个上下文：
from contextlib import ExitStack
# 动态打开多个文件
filenames = ['a.txt', 'b.txt', 'c.txt']
with ExitStack() as stack:
files = [stack.enter_context(open(f)) for f in filenames]
# 所有文件会在退出时自动关闭
for f in files:
print(f.read())异常处理
@contextmanager
def transaction(db):
"""数据库事务上下文管理器"""
try:
db.begin()
yield db
db.commit() # 正常退出时提交
except Exception:
db.rollback() # 异常时回滚
raise # 重新抛出异常
with transaction(db) as t:
t.execute("INSERT INTO users VALUES (1, 'Alice')")
t.execute("INSERT INTO orders VALUES (100, 1)")嵌套上下文
# 同时打开输入和输出文件
with open('input.txt') as fin, open('output.txt', 'w') as fout:
for line in fin:
fout.write(line.upper())
# 等价于
with open('input.txt') as fin:
with open('output.txt', 'w') as fout:
for line in fin:
fout.write(line.upper())异步上下文管理器
class AsyncDB:
"""异步数据库连接"""
async def __aenter__(self):
self.conn = await connect()
return self.conn
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.conn.close()
async with AsyncDB() as conn:
await conn.execute('SELECT 1')使用 @asynccontextmanager：
from contextlib import asynccontextmanager
@asynccontextmanager
async def async_timer(name):
"""异步计时器"""
start = time.perf_counter()
yield
elapsed = time.perf_counter() - start
print(f"{name}: {elapsed:.3f}s")
async with async_timer("async query"):
await db.query("SELECT * FROM users")常见场景
场景一：文件操作
@contextmanager
def safe_write(filepath):
"""安全写入文件：先写入临时文件，成功后再重命名"""
tmp_path = filepath + '.tmp'
try:
with open(tmp_path, 'w') as f:
yield f
os.rename(tmp_path, filepath) # 原子操作
except Exception:
if os.path.exists(tmp_path):
os.remove(tmp_path) # 失败时删除临时文件
raise
with safe_write('config.json') as f:
json.dump(config, f)
# 只有写入成功才会替换原文件场景二：临时修改环境
import os
@contextmanager
def temp_env(**kwargs):
"""临时修改环境变量"""
old_values = {}
for key, value in kwargs.items():
old_values[key] = os.environ.get(key)
os.environ[key] = value
try:
yield
finally:
for key, old_value in old_values.items():
if old_value is None:
os.environ.pop(key, None)
else:
os.environ[key] = old_value
with temp_env(DATABASE_URL="postgresql://test", DEBUG="1"):
# 在这个代码块中环境变量被临时修改
run_tests()
# 退出后环境变量恢复原值场景三：性能分析
import cProfile
import pstats
@contextmanager
def profile(output_file=None):
"""性能分析上下文管理器"""
profiler = cProfile.Profile()
profiler.enable()
try:
yield profiler
finally:
profiler.disable()
stats = pstats.Stats(profiler)
stats.sort_stats('cumulative')
if output_file:
stats.dump_stats(output_file)
else:
stats.print_stats(20) # 打印前 20 行
with profile("profile.stats"):
expensive_operation()场景四：锁管理
import threading
@contextmanager
def acquire_lock(lock, timeout=5):
"""带超时的锁管理"""
acquired = lock.acquire(timeout=timeout)
if not acquired:
raise TimeoutError("获取锁超时")
try:
yield
finally:
lock.release()
lock = threading.Lock()
with acquire_lock(lock, timeout=10):
# 在锁保护下执行操作
update_shared_resource()注意事项
__exit__返回 True 会吞掉异常，除非有明确需求，否则应返回 False- 使用
@contextmanager时，yield 之前的代码相当于__enter__，之后的代码相当于__exit__ @contextmanager中 yield 之后的清理代码应放在 try/finally 中，确保即使 yield 本身抛出异常也能清理- 异步上下文管理器使用
async with，不能用普通的with - 不要在
__exit__中执行可能抛出异常的操作，这会掩盖原始异常
进阶用法
可重入的上下文管理器
class ReentrantManager:
"""可重入的上下文管理器"""
_depth = 0
def __enter__(self):
if self._depth == 0:
print("首次进入")
self._depth += 1
return self
def __exit__(self, *exc):
self._depth -= 1
if self._depth == 0:
print("最终退出")
return False上下文管理器工厂
class DBTransaction:
"""可配置的事务管理器"""
def __init__(self, db, isolation_level='READ COMMITTED'):
self.db = db
self.isolation_level = isolation_level
def __enter__(self):
self.db.execute(f"SET TRANSACTION ISOLATION LEVEL {self.isolation_level}")
self.db.begin()
return self.db
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
self.db.commit()
else:
self.db.rollback()
return False
# 使用
with DBTransaction(db, 'SERIALIZABLE') as t:
t.execute("UPDATE accounts SET balance = balance - 100")
t.execute("UPDATE accounts SET balance = balance + 100")使用 contextlib.nullcontext
Python 3.7+ 提供了空上下文管理器，用于条件性使用上下文：
from contextlib import nullcontext
def process(lock=None):
"""可选的锁保护"""
ctx = lock if lock else nullcontext()
with ctx:
# 有锁时加锁执行，无锁时直接执行
do_work()上下文管理器与装饰器结合
from functools import wraps
def contextmanager_to_decorator(cm_factory):
"""将上下文管理器工厂转换为装饰器"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
with cm_factory():
return func(*args, **kwargs)
return wrapper
return decorator
# 使用
@contextmanager_to_decorator(timer)
def my_function():
# 函数执行时自动计时
expensive_computation()