大家都用过 with open() as f 来读写文件,但可能较少去实现自己的 context manager。今天我们就通过几个实用场景,来深入理解这个优雅的语法特性。
你一定用过:优雅处理资源管理
在 Python 中,如果不正确关闭文件句柄,可能带来严重后果:
# 错误示例
f = open('huge_file.txt')
content = f.read()
# 忘记调用 f.close()
# 潜在问题:
# 1. 文件句柄泄露:操作系统能打开的文件数是有限的
# 2. 数据丢失:写入的数据可能还在缓冲区,未真正写入磁盘
# 3. 文件锁定:其他程序可能无法访问该文件
这就是为什么我们推荐使用 with 语句:
with open('huge_file.txt') as f:
content = f.read()
# 这里自动调用了 f.close(),即使发生异常也会关闭
那么,为什么使用了 with 可以自动调用 f.close() 呢?
从一个数据分析场景说起
假设你正在处理大量临时数据文件,下载后需要及时清理以节省磁盘空间:
def process_data():
# 未使用 with 的写法
try:
data = download_large_file()
result = analyze(data)
cleanup_temp_files()
return result
except Exception as e:
cleanup_temp_files()
raise e
这种写法有几个问题:
- cleanup 逻辑重复了
- 如果中间加入 return ,容易忘记cleanup
- 代码结构不够优雅
让我们改用 context manager 的方式:
class DataManager:
def __enter__(self):
self.data = download_large_file()
return self.data
def __exit__(self, exc_type, exc_value, traceback):
cleanup_temp_files()
return False # 不吞掉异常
def process_data():
with DataManager() as data:
return analyze(data) # 自动cleanup,更简洁
如上,当我们定义了 __enter__ 和 __exit__ 方法,Python 会在使用 with 语句时自动调用 __enter__,离开 with 语句时定义的作用域时自动调用 __exit__。
__exit__方法的返回值决定了异常是否会被"吞掉"(suppressed):
- 如果 __exit__ 返回 True :
如果在上下文管理器块中发生了异常,这个异常会被抑制
程序会继续正常执行,就像没有发生异常一样
- 如果 __exit__ 返回 False 或 None (默认):
异常会被重新抛出
程序会按照正常的异常处理流程执行
常见应用场景
1. 资源管理
with open('huge_file.txt') as f:
content = f.read()
除了文件操作,还包括:
- 数据库连接
- 网络连接
- 临时文件处理
2. 代码计时器
class Timer:
def __enter__(self):
self.start = time.time() # 步骤1:进入 with 代码块时执行
return self
def __exit__(self, *args):
self.end = time.time() # 步骤3:离开 with 代码块时执行
print(f'耗时: {self.end - self.start:.2f}秒')
# 使用示例
with Timer():
time.sleep(1.5) # 步骤2:执行 with 代码块内的代码
# 步骤3会在这里自动执行,即使发生异常也会执行
3. 线程锁
from threading import Lock
class SafeCounter:
def __init__(self):
self._counter = 0
self._lock = Lock()
@property
def counter(self):
with self._lock: # 自动加锁解锁
return self._counter
@contextmanager 装饰器解析
除了定义类,还可以用装饰器 @contextmanager 来创建 context manager 。
当我们使用 @contextmanager 装饰一个生成器函数时,装饰器会:
- 创建一个新的类,实现 __enter__ 和 __exit__ 方法
- 将我们的生成器函数分成三部分:
- yield 之前的代码放入 __enter__
- yield 的值作为 __enter__ 的返回值
- yield 之后的代码放入 __exit__
例如:
import os
from contextlib import contextmanager
import time
# 方式1:使用 @contextmanager 装饰器
@contextmanager
def temp_file(filename):
# __enter__ 部分
print(f"创建临时文件: {filename}")
with open(filename, 'w') as f:
f.write('一些临时数据')
try:
yield filename # 返回值
finally:
# __exit__ 部分
print(f"清理临时文件: {filename}")
os.remove(filename)
# 方式2:使用传统的类实现
class TempFileManager:
def __init__(self, filename):
self.filename = filename
def __enter__(self):
print(f"创建临时文件: {self.filename}")
with open(self.filename, 'w') as f:
f.write('一些临时数据')
return self.filename
def __exit__(self, exc_type, exc_value, traceback):
print(f"清理临时文件: {self.filename}")
os.remove(self.filename)
return False
# 测试代码
def process_file(filepath):
print(f"处理文件: {filepath}")
time.sleep(1) # 模拟一些处理过程
if "error" in filepath:
raise ValueError("发现错误文件名!")
def test_context_manager():
print("\n1. 测试 @contextmanager 装饰器版本:")
try:
with temp_file("test1.txt") as f:
process_file(f)
print("正常完成")
except ValueError as e:
print(f"捕获到异常: {e}")
print("\n2. 测试类实现版本:")
try:
with TempFileManager("test2.txt") as f:
process_file(f)
print("正常完成")
except ValueError as e:
print(f"捕获到异常: {e}")
print("\n3. 测试异常情况:")
try:
with temp_file("error.txt") as f:
process_file(f)
print("正常完成")
except ValueError as e:
print(f"捕获到异常: {e}")
if __name__ == "__main__":
test_context_manager()
输出如下:
1. 测试 @contextmanager 装饰器版本:
创建临时文件: test1.txt
处理文件: test1.txt
清理临时文件: test1.txt
正常完成
2. 测试类实现版本:
创建临时文件: test2.txt
处理文件: test2.txt
清理临时文件: test2.txt
正常完成
3. 测试异常情况:
创建临时文件: error.txt
处理文件: error.txt
清理临时文件: error.txt
捕获到异常: 发现错误文件名!
高级用法:异常处理
__exit__ 方法可以优雅处理异常:
import sqlite3
import time
from contextlib import contextmanager
class Transaction:
def __init__(self, db_path):
self.db_path = db_path
def __enter__(self):
print("开始事务...")
self.conn = sqlite3.connect(self.db_path)
self.conn.execute('BEGIN TRANSACTION')
return self.conn
def __exit__(self, exc_type, exc_value, traceback):
if exc_type is None:
print("提交事务...")
self.conn.commit()
else:
print(f"回滚事务... 异常: {exc_type.__name__}: {exc_value}")
self.conn.rollback()
self.conn.close()
return False # 不吞掉异常
# 为了对比,我们也实现一个装饰器版本
@contextmanager
def transaction(db_path):
print("开始事务...")
conn = sqlite3.connect(db_path)
conn.execute('BEGIN TRANSACTION')
try:
yield conn
print("提交事务...")
conn.commit()
except Exception as e:
print(f"回滚事务... 异常: {type(e).__name__}: {e}")
conn.rollback()
raise # 重新抛出异常
finally:
conn.close()
def init_db(db_path):
"""初始化数据库"""
conn = sqlite3.connect(db_path)
conn.execute('''
CREATE TABLE IF NOT EXISTS accounts (
id INTEGER PRIMARY KEY,
name TEXT,
balance REAL
)
''')
# 插入初始数据
conn.execute('DELETE FROM accounts') # 清空旧数据
conn.execute('INSERT INTO accounts (name, balance) VALUES (?, ?)', ('Alice', 1000))
conn.execute('INSERT INTO accounts (name, balance) VALUES (?, ?)', ('Bob', 1000))
conn.commit()
conn.close()
def transfer_money(conn, from_name, to_name, amount):
"""转账操作"""
print(f"转账: {from_name} -> {to_name}, 金额: {amount}")
# 模拟一些延迟,便于观察
time.sleep(1)
# 扣款
cursor = conn.execute(
'UPDATE accounts SET balance = balance - ? WHERE name = ? AND balance >= ?',
(amount, from_name, amount)
)
if cursor.rowcount == 0:
raise ValueError(f"{from_name} 余额不足或账户不存在!")
# 模拟可能的错误情况
if to_name == "ErrorUser":
raise ValueError("目标账户不存在!")
# 入账
conn.execute(
'UPDATE accounts SET balance = balance + ? WHERE name = ?',
(amount, to_name)
)
def show_balances(db_path):
"""显示所有账户余额"""
conn = sqlite3.connect(db_path)
cursor = conn.execute('SELECT name, balance FROM accounts')
print("\n当前余额:")
for name, balance in cursor:
print(f"{name}: {balance}")
conn.close()
def test_transactions():
db_path = "test_transactions.db"
init_db(db_path)
print("\n1. 测试正常转账:")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "Bob", 300)
print("转账成功!")
except Exception as e:
print(f"转账失败: {e}")
show_balances(db_path)
print("\n2. 测试余额不足:")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "Bob", 2000)
print("转账成功!")
except Exception as e:
print(f"转账失败: {e}")
show_balances(db_path)
print("\n3. 测试无效账户:")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "ErrorUser", 100)
print("转账成功!")
except Exception as e:
print(f"转账失败: {e}")
show_balances(db_path)
print("\n4. 使用装饰器版本测试:")
try:
with transaction(db_path) as conn:
transfer_money(conn, "Bob", "Alice", 200)
print("转账成功!")
except Exception as e:
print(f"转账失败: {e}")
show_balances(db_path)
if __name__ == "__main__":
test_transactions()
输出如下:
1. 测试正常转账:
开始事务...
转账: Alice -> Bob, 金额: 300
提交事务...
转账成功!
当前余额:
Alice: 700.0
Bob: 1300.0
2. 测试余额不足:
开始事务...
转账: Alice -> Bob, 金额: 2000
回滚事务... 异常: ValueError: Alice 余额不足或账户不存在!
转账失败: Alice 余额不足或账户不存在!
当前余额:
Alice: 700.0
Bob: 1300.0
3. 测试无效账户:
开始事务...
转账: Alice -> ErrorUser, 金额: 100
回滚事务... 异常: ValueError: 目标账户不存在!
转账失败: 目标账户不存在!
当前余额:
Alice: 700.0
Bob: 1300.0
4. 使用装饰器版本测试:
开始事务...
转账: Bob -> Alice, 金额: 200
提交事务...
转账成功!
当前余额:
Alice: 900.0
Bob: 1100.0
实用建议
- 及时清理:__exit__ 确保资源释放
- 异常透明:通常返回 False,让异常继续传播
- 功能单一:一个 context manager 只做一件事
- 考虑可组合:多个 with 可以组合使用
小结
with 语句是 Python 中非常优雅的特性,善用它可以:
- 自动管理资源
- 简化异常处理
- 提高代码可读性
建议大家在处理需要配对操作的场景(开启/关闭、加锁/解锁、创建/删除等)时,优先考虑使用 with 语句。