7 minutes
上下文管理器与 with 语句
上下文管理器(Context Manager)是 Python 中用于资源管理的核心机制。无论你是在读写文件、获取锁、打开数据库连接还是管理网络会话——with 语句都能确保资源被正确释放,即使发生异常也不例外。
为什么需要上下文管理器
传统方式的痛点
来看一个常见的文件操作:
# 传统方式:容易忘记关闭文件
f = open("data.txt", "r")
content = f.read()
f.close() # 如果前面的代码抛异常,这一行不会执行!
# 改进:用 try/finally 确保关闭
f = open("data.txt", "r")
try:
content = f.read()
finally:
f.close() # 无论是否异常,都会执行
虽然 try/finally 能解决问题,但代码显得冗长。如果每次资源操作都要写 try/finally,不仅繁琐而且容易遗漏。
使用 with 语句
# with 语句 —— 简洁且安全
with open("data.txt", "r") as f:
content = f.read()
# 离开 with 块后,文件自动关闭
# 等价于 try/finally,但更简洁
with 语句的工作原理
with 语句的本质是执行上下文管理器的协议:
with EXPRESSION as VARIABLE:
BLOCK
等价于:
manager = EXPRESSION
value = manager.__enter__()
VARIABLE = value
try:
BLOCK
except Exception as e:
if not manager.__exit__(type(e), e, e.__traceback__):
raise
else:
manager.__exit__(None, None, None)
实现上下文管理器
方式一:类实现 __enter__ 和 __exit__
class ManagedFile:
def __init__(self, filename: str, mode: str = "r"):
self.filename = filename
self.mode = mode
self.file = None
def __enter__(self):
print(f"打开文件: {self.filename}")
self.file = open(self.filename, self.mode)
return self.file # 返回给 as 子句
def __exit__(self, exc_type, exc_val, exc_tb):
if self.file:
self.file.close()
print(f"关闭文件: {self.filename}")
# 返回 True 会抑制异常,返回 False(默认)则异常继续传播
if exc_type is not None:
print(f"发生异常: {exc_val}")
return False # 不抑制异常
# 正常使用
with ManagedFile("test.txt", "w") as f:
f.write("Hello, World!")
# 输出:
# 打开文件: test.txt
# 关闭文件: test.txt
# 异常时的使用
try:
with ManagedFile("test.txt", "r") as f:
content = f.read()
raise ValueError("模拟错误")
except ValueError as e:
print(f"捕获到异常: {e}")
# 输出:
# 打开文件: test.txt
# 发生异常: 模拟错误
# 关闭文件: test.txt
# 捕获到异常: 模拟错误
方式二:@contextmanager 装饰器
更简洁的方式是使用 contextlib 模块的 @contextmanager 装饰器:
from contextlib import contextmanager
@contextmanager
def managed_file(filename: str, mode: str = "r"):
"""使用生成器实现的上下文管理器"""
print(f"打开文件: {filename}")
file = open(filename, mode)
try:
yield file # yield 之前的代码相当于 __enter__
finally:
file.close()
print(f"关闭文件: {filename}")
with managed_file("test.txt", "w") as f:
f.write("Hello from generator!")
@contextmanager 的异常处理
@contextmanager
def divider(a: float, b: float):
"""安全的除法上下文管理器"""
print(f"准备计算: {a} / {b}")
try:
yield a / b
except ZeroDivisionError as e:
print(f"计算错误: {e}")
yield float("inf") # 返回无穷大作为替代值
finally:
print("计算完成")
with divider(10, 2) as result:
print(f"结果: {result}") # 结果: 5.0
with divider(10, 0) as result:
print(f"结果: {result}") # 结果: inf
contextlib 实用工具
Python 的 contextlib 模块提供了丰富的上下文管理器工具:
closing
确保对象被正确关闭(调用 close() 方法):
from contextlib import closing
import urllib.request
with closing(urllib.request.urlopen("https://httpbin.org/get")) as response:
data = response.read()
print(f"状态码: {response.status}")
# 离开 with 块后,response.close() 自动调用
suppress
忽略指定的异常:
from contextlib import suppress
import os
# 传统方式
try:
os.remove("temp.txt")
except FileNotFoundError:
pass
# 使用 suppress
with suppress(FileNotFoundError):
os.remove("temp.txt")
# 可以同时忽略多种异常
with suppress(FileNotFoundError, PermissionError):
os.remove("temp.txt")
nullcontext
需要一个上下文管理器但什么都不做时使用:
from contextlib import nullcontext
from typing import ContextManager
def process_data(data: str, use_lock: bool = False) -> None:
"""根据参数决定是否需要加锁"""
lock = threading.Lock() if use_lock else nullcontext()
with lock: # 有锁时加锁,没有锁时 nullcontext 什么也不做
print(f"Processing: {data}")
redirect_stdout 和 redirect_stderr
临时重定向标准输出/错误:
from contextlib import redirect_stdout, redirect_stderr
import io
def noisy_function():
print("This goes to stdout")
print("This goes to stderr", file=sys.stderr)
# 捕获输出到字符串
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()
with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
noisy_function()
print(f"捕获的 stdout: {stdout_capture.getvalue()}")
print(f"捕获的 stderr: {stderr_capture.getvalue()}")
# 恢复后,输出回到正常位置
noisy_function()
ExitStack
动态管理多个上下文管理器:
from contextlib import ExitStack
def open_many_files(filenames: list[str]) -> list:
"""打开多个文件,确保所有文件都能被正确关闭"""
files = []
with ExitStack() as stack:
for filename in filenames:
f = open(filename, "w")
stack.callback(f.close) # 注册清理回调
files.append(f)
# 返回前将文件列表与 ExitStack 关联
return files # 注意:这里文件会在 with 块退出时被关闭!
# 正确做法:让调用者也使用 with
def process_files(filenames: list[str]) -> None:
with ExitStack() as stack:
files = [
stack.enter_context(open(fname))
for fname in filenames
]
# 在 with 块内安全使用所有文件
for f, name in zip(files, filenames):
f.write(f"Content for {name}")
嵌套上下文管理器
多个 with 的嵌套
# 嵌套写法
with open("input.txt") as infile:
with open("output.txt", "w") as outfile:
outfile.write(infile.read())
# 等效的扁平写法(Python 3.1+)
with open("input.txt") as infile, open("output.txt", "w") as outfile:
outfile.write(infile.read())
数据库事务示例
import sqlite3
from contextlib import contextmanager
@contextmanager
def database_transaction(db_path: str):
"""数据库事务上下文管理器"""
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
try:
yield cursor
conn.commit()
print("事务提交成功")
except Exception as e:
conn.rollback()
print(f"事务回滚: {e}")
raise
finally:
conn.close()
# 嵌套上下文管理器
def transfer_money(from_id: int, to_id: int, amount: float) -> None:
with database_transaction("bank.db") as cursor:
cursor.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?",
(amount, from_id))
cursor.execute("UPDATE accounts SET balance = balance + ? WHERE id = ?",
(amount, to_id))
自定义上下文管理器实战
实战 1:性能计时器
import time
from contextlib import contextmanager
@contextmanager
def timer(name: str = "block"):
"""代码块执行计时器"""
start = time.perf_counter()
print(f"[{name}] 开始执行...")
try:
yield
finally:
elapsed = time.perf_counter() - start
print(f"[{name}] 耗时: {elapsed:.4f}s")
with timer("数据处理"):
time.sleep(0.5)
result = sum(i * i for i in range(10_000_000))
print(f"计算结果: {result}")
# 输出:
# [数据处理] 开始执行...
# 计算结果: 333333283333335000000
# [数据处理] 耗时: 0.5234s
实战 2:临时目录切换
import os
from pathlib import Path
from contextlib import contextmanager
@contextmanager
def change_directory(target_dir: str):
"""临时切换工作目录"""
original_dir = os.getcwd()
try:
os.chdir(target_dir)
print(f"切换到目录: {target_dir}")
yield
finally:
os.chdir(original_dir)
print(f"恢复目录: {original_dir}")
# 使用示例
with change_directory(Path.home() / "Documents"):
print(f"当前目录: {os.getcwd()}")
# 在这里做文件操作...
print(f"当前目录: {os.getcwd()}") # 已恢复
实战 3:线程锁上下文
import threading
from contextlib import contextmanager
class SharedCounter:
"""线程安全的计数器"""
def __init__(self):
self.value = 0
self._lock = threading.Lock()
def increment(self) -> None:
with self._lock: # 锁的上下文管理
self.value += 1
def decrement(self) -> None:
with self._lock:
self.value -= 1
@contextmanager
def batch_update(self):
"""批量更新上下文(一次加锁,多次操作)"""
with self._lock:
print("开始批量更新")
yield
print("批量更新结束")
counter = SharedCounter()
# 使用锁上下文
def worker():
for _ in range(10000):
counter.increment()
threads = [threading.Thread(target=worker) for _ in range(10)]
for t in threads:
t.start()
for t in threads:
t.join()
print(f"最终结果 (期望 100000): {counter.value}")
# 批量更新
with counter.batch_update():
counter.value += 10
counter.value += 20
counter.value += 30
实战 4:临时环境变量
import os
from contextlib import contextmanager
@contextmanager
def set_env(**environ: str):
"""临时设置环境变量"""
old_env = {}
try:
for key, value in environ.items():
old_env[key] = os.environ.get(key)
os.environ[key] = value
print(f"设置 {key}={value}")
yield
finally:
for key in environ:
if old_env[key] is None:
del os.environ[key]
else:
os.environ[key] = old_env[key]
print(f"恢复 {key}={old_env[key]}")
with set_env(DATABASE_URL="sqlite:///test.db", DEBUG="true"):
print(f"DATABASE_URL: {os.environ['DATABASE_URL']}")
print(f"DEBUG: {os.environ['DEBUG']}")
print(f"DATABASE_URL: {os.environ.get('DATABASE_URL')}") # 已恢复
实战 5:连接池管理
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
@dataclass
class Connection:
"""模拟数据库连接"""
id: int
created_at: datetime = datetime.now()
def query(self, sql: str) -> str:
return f"[连接 #{self.id}] 执行: {sql}"
def close(self) -> None:
print(f"关闭连接 #{self.id}")
class ConnectionPool:
"""数据库连接池"""
def __init__(self, min_size: int = 2, max_size: int = 10):
self.min_size = min_size
self.max_size = max_size
self._idle: list[Connection] = [Connection(i) for i in range(min_size)]
self._active: set[Connection] = set()
self._counter = min_size
@contextmanager
def get_connection(self) -> Connection:
"""获取连接(使用后自动归还)"""
conn = self._acquire()
try:
yield conn
finally:
self._release(conn)
def _acquire(self) -> Connection:
if self._idle:
conn = self._idle.pop()
self._active.add(conn)
print(f"从连接池获取连接 #{conn.id}")
return conn
if len(self._active) < self.max_size:
self._counter += 1
conn = Connection(self._counter)
self._active.add(conn)
print(f"创建新连接 #{conn.id}")
return conn
raise RuntimeError("连接池已满")
def _release(self, conn: Connection) -> None:
self._active.discard(conn)
if len(self._idle) < self.max_size:
self._idle.append(conn)
print(f"归还连接 #{conn.id} 到连接池")
else:
conn.close()
pool = ConnectionPool(min_size=2)
# 使用连接
with pool.get_connection() as conn:
result = conn.query("SELECT * FROM users")
print(result)
with pool.get_connection() as conn:
result = conn.query("SELECT * FROM orders")
print(result)
类实现 vs 装饰器实现
两种实现上下文管理器的方式各有优劣:
| 对比项 | 类实现 | @contextmanager |
|---|---|---|
| 复杂度 | 稍复杂 | 简洁 |
| 异常处理 | 灵活(通过返回值控制抑制) | 有限(只能通过 try/except) |
| 需要记住状态 | 适合(self 可以存储状态) | 不太适合 |
| 可读性 | 结构清晰 | 非常简洁 |
| 需要单独的类 | 是 | 否(函数即可) |
选择原则
# 需要状态管理 → 类实现
class RetryContext:
def __init__(self, max_retries: int = 3):
self.max_retries = max_retries
self.attempt = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is not None and self.attempt < self.max_retries:
self.attempt += 1
print(f"重试 {self.attempt}/{self.max_retries}")
return True # 抑制异常,继续重试
return False
# 简单场景 → @contextmanager
@contextmanager
def simple_timer():
start = time.perf_counter()
yield
print(f"耗时: {time.perf_counter() - start:.4f}s")
__enter__ 返回值详解
__enter__ 的返回值通过 as 子句传递给变量:
class Connection:
def __enter__(self):
print("建立连接")
return self # 返回自身
def __exit__(self, *args):
print("关闭连接")
def query(self, sql: str):
return f"查询: {sql}"
class Config:
def __enter__(self):
print("加载配置")
return {"debug": True, "port": 8080} # 返回字典
def __exit__(self, *args):
print("保存配置")
# conn 是 Connection 实例
with Connection() as conn:
print(conn.query("SELECT 1"))
# config 是字典
with Config() as config:
print(config["debug"])
常见陷阱
陷阱 1:__exit__ 中返回了 True 但不应该
class SilentError:
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
print(f"捕获到: {exc_val}")
return True # 抑制了异常!
with SilentError() as se:
raise ValueError("严重错误!")
# 不会触发异常!这通常不是你想要的
陷阱 2:@contextmanager 生成器中没有 try/finally
@contextmanager
def bad_context():
print("进入")
yield "value"
# 如果 yield 之前的代码或 with 块中抛异常,
# yield 之后的代码不会执行!
print("退出") # 异常时不会执行!
@contextmanager
def good_context():
print("进入")
try:
yield "value"
finally:
print("退出") # 无论如何都会执行
with good_context() as v:
print(v)
raise ValueError("测试")
# 输出:
# 进入
# value
# 退出
陷阱 3:同时管理多个资源时忘记处理顺序
# 正确:先打开的资源后关闭
with open("a.txt") as a, open("b.txt") as b:
pass # a 先打开,b 后打开;退出时 b 先关闭,a 后关闭
# 等价于:
with open("a.txt") as a:
with open("b.txt") as b:
pass
小结
上下文管理器是 Python 中优雅管理资源的关键机制。本篇涵盖了:
- with 语句:Python 推荐的安全资源管理方式
__enter__/__exit__协议:类实现上下文管理器的基础@contextmanager装饰器:用生成器函数简化上下文管理器实现- contextlib 工具集:
closing、suppress、nullcontext、redirect_stdout、ExitStack - 嵌套上下文管理器:同时管理多个资源
- 实战示例:计时器、目录切换、线程锁、环境变量、连接池
何时使用上下文管理器:
- 文件操作
- 数据库连接和事务
- 线程锁
- 网络连接
- 临时环境修改
- 计时和性能分析
- 任何需要"进入-退出"模式的操作
掌握上下文管理器后,你就能写出更安全、更 Pythonic 的代码。
Summary: with 语句、enter/exit、@contextmanager、contextlib