Python with 語(yǔ)句的深入理解：優(yōu)雅處理資源管理 @contextmanager
作者：piperliu
大家都用過(guò) with open() as f 來(lái)讀寫(xiě)文件，但可能較少去實(shí)現(xiàn)自己的 context manager。今天我們就通過(guò)幾個(gè)實(shí)用場(chǎng)景，來(lái)深入理解這個(gè)優(yōu)雅的語(yǔ)法特性。
大家都用過(guò) with open() as f 來(lái)讀寫(xiě)文件，但可能較少去實(shí)現(xiàn)自己的 context manager。今天我們就通過(guò)幾個(gè)實(shí)用場(chǎng)景，來(lái)深入理解這個(gè)優(yōu)雅的語(yǔ)法特性。
你一定用過(guò)：優(yōu)雅處理資源管理
在 Python 中，如果不正確關(guān)閉文件句柄，可能帶來(lái)嚴(yán)重后果：
# 錯(cuò)誤示例
f = open('huge_file.txt')
content = f.read()
# 忘記調(diào)用 f.close()
# 潛在問(wèn)題：
# 1. 文件句柄泄露：操作系統(tǒng)能打開(kāi)的文件數(shù)是有限的
# 2. 數(shù)據(jù)丟失：寫(xiě)入的數(shù)據(jù)可能還在緩沖區(qū)，未真正寫(xiě)入磁盤(pán)
# 3. 文件鎖定：其他程序可能無(wú)法訪問(wèn)該文件這就是為什么我們推薦使用 with 語(yǔ)句：
with open('huge_file.txt') as f:
content = f.read()
# 這里自動(dòng)調(diào)用了 f.close()，即使發(fā)生異常也會(huì)關(guān)閉那么，為什么使用了 with 可以自動(dòng)調(diào)用 f.close() 呢？
從一個(gè)數(shù)據(jù)分析場(chǎng)景說(shuō)起
假設(shè)你正在處理大量臨時(shí)數(shù)據(jù)文件，下載后需要及時(shí)清理以節(jié)省磁盤(pán)空間：
def process_data():
# 未使用 with 的寫(xiě)法
try:
data = download_large_file()
result = analyze(data)
cleanup_temp_files()
return result
except Exception as e:
cleanup_temp_files()
raise e這種寫(xiě)法有幾個(gè)問(wèn)題：
- cleanup 邏輯重復(fù)了
- 如果中間加入 return ，容易忘記cleanup
- 代碼結(jié)構(gòu)不夠優(yōu)雅
讓我們改用 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) # 自動(dòng)cleanup，更簡(jiǎn)潔如上，當(dāng)我們定義了 __enter__ 和 __exit__ 方法，Python 會(huì)在使用 with 語(yǔ)句時(shí)自動(dòng)調(diào)用 __enter__，離開(kāi) with 語(yǔ)句時(shí)定義的作用域時(shí)自動(dòng)調(diào)用 __exit__。
__exit__方法的返回值決定了異常是否會(huì)被"吞掉"（suppressed）：
- 如果 __exit__ 返回 True ：
如果在上下文管理器塊中發(fā)生了異常，這個(gè)異常會(huì)被抑制
程序會(huì)繼續(xù)正常執(zhí)行，就像沒(méi)有發(fā)生異常一樣
- 如果 __exit__ 返回 False 或 None （默認(rèn)）：
異常會(huì)被重新拋出
程序會(huì)按照正常的異常處理流程執(zhí)行
常見(jiàn)應(yīng)用場(chǎng)景
1. 資源管理
with open('huge_file.txt') as f:
content = f.read()除了文件操作，還包括:
- 數(shù)據(jù)庫(kù)連接
- 網(wǎng)絡(luò)連接
- 臨時(shí)文件處理
2. 代碼計(jì)時(shí)器
class Timer:
def __enter__(self):
self.start = time.time() # 步驟1：進(jìn)入 with 代碼塊時(shí)執(zhí)行
return self
def __exit__(self, *args):
self.end = time.time() # 步驟3：離開(kāi) with 代碼塊時(shí)執(zhí)行
print(f'耗時(shí): {self.end - self.start:.2f}秒')
# 使用示例
with Timer():
time.sleep(1.5) # 步驟2：執(zhí)行 with 代碼塊內(nèi)的代碼
# 步驟3會(huì)在這里自動(dòng)執(zhí)行，即使發(fā)生異常也會(huì)執(zhí)行3. 線程鎖
from threading import Lock
class SafeCounter:
def __init__(self):
self._counter = 0
self._lock = Lock()
@property
def counter(self):
with self._lock: # 自動(dòng)加鎖解鎖
return self._counter@contextmanager 裝飾器解析
除了定義類(lèi)，還可以用裝飾器 @contextmanager 來(lái)創(chuàng)建 context manager 。
當(dāng)我們使用 @contextmanager 裝飾一個(gè)生成器函數(shù)時(shí)，裝飾器會(huì)：
- 創(chuàng)建一個(gè)新的類(lèi)，實(shí)現(xiàn) __enter__ 和 __exit__ 方法
- 將我們的生成器函數(shù)分成三部分：
- yield 之前的代碼放入 __enter__
- yield 的值作為 __enter__ 的返回值
- yield 之后的代碼放入 __exit__
例如：
import os
from contextlib import contextmanager
import time
# 方式1：使用 @contextmanager 裝飾器
@contextmanager
def temp_file(filename):
# __enter__ 部分
print(f"創(chuàng)建臨時(shí)文件: {filename}")
with open(filename, 'w') as f:
f.write('一些臨時(shí)數(shù)據(jù)')
try:
yield filename # 返回值
finally:
# __exit__ 部分
print(f"清理臨時(shí)文件: {filename}")
os.remove(filename)
# 方式2：使用傳統(tǒng)的類(lèi)實(shí)現(xiàn)
class TempFileManager:
def __init__(self, filename):
self.filename = filename
def __enter__(self):
print(f"創(chuàng)建臨時(shí)文件: {self.filename}")
with open(self.filename, 'w') as f:
f.write('一些臨時(shí)數(shù)據(jù)')
return self.filename
def __exit__(self, exc_type, exc_value, traceback):
print(f"清理臨時(shí)文件: {self.filename}")
os.remove(self.filename)
return False
# 測(cè)試代碼
def process_file(filepath):
print(f"處理文件: {filepath}")
time.sleep(1) # 模擬一些處理過(guò)程
if "error" in filepath:
raise ValueError("發(fā)現(xiàn)錯(cuò)誤文件名！")
def test_context_manager():
print("\n1. 測(cè)試 @contextmanager 裝飾器版本:")
try:
with temp_file("test1.txt") as f:
process_file(f)
print("正常完成")
except ValueError as e:
print(f"捕獲到異常: {e}")
print("\n2. 測(cè)試類(lèi)實(shí)現(xiàn)版本:")
try:
with TempFileManager("test2.txt") as f:
process_file(f)
print("正常完成")
except ValueError as e:
print(f"捕獲到異常: {e}")
print("\n3. 測(cè)試異常情況:")
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. 測(cè)試 @contextmanager 裝飾器版本:
創(chuàng)建臨時(shí)文件: test1.txt
處理文件: test1.txt
清理臨時(shí)文件: test1.txt
正常完成
2. 測(cè)試類(lèi)實(shí)現(xiàn)版本:
創(chuàng)建臨時(shí)文件: test2.txt
處理文件: test2.txt
清理臨時(shí)文件: test2.txt
正常完成
3. 測(cè)試異常情況:
創(chuàng)建臨時(shí)文件: error.txt
處理文件: error.txt
清理臨時(shí)文件: error.txt
捕獲到異常: 發(fā)現(xiàn)錯(cuò)誤文件名！高級(jí)用法：異常處理
__exit__ 方法可以?xún)?yōu)雅處理異常：
import sqlite3
import time
from contextlib import contextmanager
class Transaction:
def __init__(self, db_path):
self.db_path = db_path
def __enter__(self):
print("開(kāi)始事務(wù)...")
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("提交事務(wù)...")
self.conn.commit()
else:
print(f"回滾事務(wù)... 異常: {exc_type.__name__}: {exc_value}")
self.conn.rollback()
self.conn.close()
return False # 不吞掉異常
# 為了對(duì)比，我們也實(shí)現(xiàn)一個(gè)裝飾器版本
@contextmanager
def transaction(db_path):
print("開(kāi)始事務(wù)...")
conn = sqlite3.connect(db_path)
conn.execute('BEGIN TRANSACTION')
try:
yield conn
print("提交事務(wù)...")
conn.commit()
except Exception as e:
print(f"回滾事務(wù)... 異常: {type(e).__name__}: {e}")
conn.rollback()
raise # 重新拋出異常
finally:
conn.close()
def init_db(db_path):
"""初始化數(shù)據(jù)庫(kù)"""
conn = sqlite3.connect(db_path)
conn.execute('''
CREATE TABLE IF NOT EXISTS accounts (
id INTEGER PRIMARY KEY,
name TEXT,
balance REAL
)
''')
# 插入初始數(shù)據(jù)
conn.execute('DELETE FROM accounts') # 清空舊數(shù)據(jù)
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):
"""轉(zhuǎn)賬操作"""
print(f"轉(zhuǎn)賬: {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} 余額不足或賬戶(hù)不存在！")
# 模擬可能的錯(cuò)誤情況
if to_name == "ErrorUser":
raise ValueError("目標(biāo)賬戶(hù)不存在！")
# 入賬
conn.execute(
'UPDATE accounts SET balance = balance + ? WHERE name = ?',
(amount, to_name)
)
def show_balances(db_path):
"""顯示所有賬戶(hù)余額"""
conn = sqlite3.connect(db_path)
cursor = conn.execute('SELECT name, balance FROM accounts')
print("\n當(dāng)前余額:")
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. 測(cè)試正常轉(zhuǎn)賬:")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "Bob", 300)
print("轉(zhuǎn)賬成功！")
except Exception as e:
print(f"轉(zhuǎn)賬失敗: {e}")
show_balances(db_path)
print("\n2. 測(cè)試余額不足:")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "Bob", 2000)
print("轉(zhuǎn)賬成功！")
except Exception as e:
print(f"轉(zhuǎn)賬失敗: {e}")
show_balances(db_path)
print("\n3. 測(cè)試無(wú)效賬戶(hù):")
try:
with Transaction(db_path) as conn:
transfer_money(conn, "Alice", "ErrorUser", 100)
print("轉(zhuǎn)賬成功！")
except Exception as e:
print(f"轉(zhuǎn)賬失敗: {e}")
show_balances(db_path)
print("\n4. 使用裝飾器版本測(cè)試:")
try:
with transaction(db_path) as conn:
transfer_money(conn, "Bob", "Alice", 200)
print("轉(zhuǎn)賬成功！")
except Exception as e:
print(f"轉(zhuǎn)賬失敗: {e}")
show_balances(db_path)
if __name__ == "__main__":
test_transactions()輸出如下：
1. 測(cè)試正常轉(zhuǎn)賬:
開(kāi)始事務(wù)...
轉(zhuǎn)賬: Alice -> Bob, 金額: 300
提交事務(wù)...
轉(zhuǎn)賬成功！
當(dāng)前余額:
Alice: 700.0
Bob: 1300.0
2. 測(cè)試余額不足:
開(kāi)始事務(wù)...
轉(zhuǎn)賬: Alice -> Bob, 金額: 2000
回滾事務(wù)... 異常: ValueError: Alice 余額不足或賬戶(hù)不存在！
轉(zhuǎn)賬失敗: Alice 余額不足或賬戶(hù)不存在！
當(dāng)前余額:
Alice: 700.0
Bob: 1300.0
3. 測(cè)試無(wú)效賬戶(hù):
開(kāi)始事務(wù)...
轉(zhuǎn)賬: Alice -> ErrorUser, 金額: 100
回滾事務(wù)... 異常: ValueError: 目標(biāo)賬戶(hù)不存在！
轉(zhuǎn)賬失敗: 目標(biāo)賬戶(hù)不存在！
當(dāng)前余額:
Alice: 700.0
Bob: 1300.0
4. 使用裝飾器版本測(cè)試:
開(kāi)始事務(wù)...
轉(zhuǎn)賬: Bob -> Alice, 金額: 200
提交事務(wù)...
轉(zhuǎn)賬成功！
當(dāng)前余額:
Alice: 900.0
Bob: 1100.0實(shí)用建議
- 及時(shí)清理：__exit__ 確保資源釋放
- 異常透明：通常返回 False，讓異常繼續(xù)傳播
- 功能單一：一個(gè) context manager 只做一件事
- 考慮可組合：多個(gè) with 可以組合使用
小結(jié)
with 語(yǔ)句是 Python 中非常優(yōu)雅的特性，善用它可以:
- 自動(dòng)管理資源
- 簡(jiǎn)化異常處理
- 提高代碼可讀性
建議大家在處理需要配對(duì)操作的場(chǎng)景（開(kāi)啟/關(guān)閉、加鎖/解鎖、創(chuàng)建/刪除等）時(shí)，優(yōu)先考慮使用 with 語(yǔ)句。
責(zé)任編輯：武曉燕
來(lái)源：
Piper蛋窩
