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