Python with 語句的深入理解:優雅處理資源管理 @contextmanager
大家都用過 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 語句。