Last Updated on 2024-09-11 by Clay
Python 中的上下文管理器裝飾器 @contextmanager 是 contextlib 模組中的一個裝飾器，可以讓開發人員簡便地創造屬於我們的上下文管理器（Context Manager）。
什麼是上下文管理器呢？簡單來說，就是定義一個處理某種既定邏輯的區塊，在『進入』和『退出』這個區塊時自動執行某些程式碼。在 Python 中最常見到的辨識搭配使用 with 語法，來自動處理文件開啟/關閉、資源分配/釋放等操作。
沒有使用 @contextmanager 的上下文管理器
假如我們要建立一個管理文件的上下文管理器，原先的作法是定義一個類別，然後定義 __enter__() 和 __exit__() 來控制進入和退出時的操作邏輯。
from typing import Any
class FileManager:
def __init__(self, filename: str, mode: str) -> None:
self.file = None
self.filename = filename
self.mode = mode
def __enter__(self) -> Any:
print("Open file.")
self.file = open(file=self.filename, mode=self.mode)
return self.file
def __exit__(self, exc_type, exc_value, traceback) -> None:
print("Close file.")
if self.file:
self.file.close()
if __name__ == "__main__":
with FileManager("test.log", "w") as f:
f.write("This is a test log.")
Output:
Open file.
Close file.
當然我們可以檢查一下，應該有成功寫入文件。總之，我們定義的類別 FileManager 會在使用 with 語法進入時打開文件並回傳文件物件 self.file，並且在 with 定義的段落結束時，自動關閉文件。
這是一個有效、傳統的寫法，下面我們來看看 @contextmanager 簡潔的寫法。
使用 @contextmanager
@contextmanager 的目的就在於簡化前一小截我們定義的方式。
from contextlib import contextmanager
@contextmanager
def open_file(filename: str, mode: str):
print("Open file.")
f = open(file=filename, mode=mode)
try:
yield f
finally:
print("Close file.")
f.close
if __name__ == "__main__":
with open_file("test.log", "r") as f:
print(f.read())
Output:
Open file.
This is a test log.
Close file.
程式碼所作的事情與前述的傳統類別定義做的事情一模一樣，就是簡化了語法。簡單來說，@contextmanager 其實建立了一個生成器。這個生成器在遇到 yield 時會暫停，並把 yield 的值返回給 with 語法後的變數（在我們的範例程式碼中就是 f）。而退出時，上下文管理器無論如何都會執行 yield 後的程式碼，所以非常適合接 try... finally 的語法。
以下幾個場景非常適合上下文管理器：
- 釋放資源：在處理過程後一定要釋放出資源的邏輯，適合使用上下文管理器的自動退出
- 臨時改變變數值：假設有一應用場景，需要操作時改變變數值，操作結束後恢復原狀，則可以使用上下文管理器（我最近遇到的一個這樣的應用場景是更新資料庫，更新過程中我的系統得處於 maintain 的狀態，可是有時使用者上傳的資料格式是錯誤的... 但無論如何，我終究還是得恢復 online 的狀態）
- 鎖（lock）的存取與釋放
總結以上幾點，我們可以理解 @contextmanager 的使用方法是：
@contextmanager用來創建簡化的上下文管理器，通過yield來分別定義進入和退出上下文時要做的事。- 進入上下文時，執行
yield之前的代碼；退出上下文時，執行yield後的代碼，通常是清理工作。 - 搭配
try和finally能確保即使發生異常也能正確釋放資源。
References
- Python Docs - contextlib — Utilities for with-statement contexts
- Context Managers and Python's with Statement