Python 中的 __enter__、__exit__與 contextmanager 教學

Jesshsieh
5 min readFeb 5, 2024

--

情境管理器在 Python 中常用來做資源管理,以下是一些常見的應用場景:

  • 執行完特定任務後關閉文件
  • 與資料庫連線完成操作後關閉連線
  • 在完成 A 任務後必須執行 B 任務的固定流程中,確保流程的順利進行。

不論是將檔案關閉或是關閉與資料庫連線都可以防止資源的浪費,減少對效能的影響,在這種時候我們通常會使用情境管理器去管理資源。

使用情境管理器的優點如下:

  • 維持程式碼簡潔、重用性
  • 簡化資源管理,不需要人工關閉資源
  • 正常關閉文件、連線,可以避免資源洩漏

以上簡單介紹完情境管理器,我們來看一些例子。

最基本的檔案操作如下:

file = open("test.txt", "w")
file.write("test123456")
file.close()

這個方法的壞處是,若執行的過程中發生錯誤,有可能無法正常關閉檔案,改善過後的寫法為:

try:
file = open("test.txt", "w")
file.write("test123456")
finally:
file.close()

儘管這種方法已經確保發生錯誤時可以正常關閉檔案,但還需要手動將檔案關閉,較為繁瑣。

Python 提供了更優雅的方法處理這個問題,
也就是透過 __enter__ 與 __exit__ 兩個 Dunder method 去管理資源。

接下來實作一個用於檔案管理的類別

class FileOpen:
def __init__(self, filename, task):
self.filename = filename
self.task = task
def __enter__(self):
print("writing")
self.file = open(self.filename, self.task)
self.file.write("test123456")
return self.file
def __exit__(self, exc_type, exc_val, exc_tb):
self.file.close()
print("file has been closed")

with FileOpen("test.txt", "w") as f:
print("file has been written.")
print(f.closed)

這是一個寫成 class 的情境管理器,可以直接從 with語句去理解。

  1. 首先執行到 with 語句時,會先實例化 FileOpenclass,filenametask 的值會傳入 __init__ 方法,這兩個參數分別定義了檔案名稱與任務
  2. 接著 __enter__ 方法會被調用,執行打開檔案並且寫入內容的任務,在 __enter__return self.file,此時 self.file 會被 with 語句後方的 as f 所接住,並且執行 print("file has been written.")
  3. 在離開 with 語句區塊後,__exit__ 方法會被觸發,因此會將檔案關閉並且執行 print("file has been closed")
  4. 最後則是多加一個判斷,確認檔案已經被關閉。

實際執行的結果如下:
>>>writing
>>>file has been written.
>>>file has been closed
>>>True

可以發現檔案確實有被關閉,解決了我們的問題。
另外一提,即使 with 語句內出錯了,__exit__() 方法仍會被呼叫,還是可以確保檔案在任何情況下會被正常關閉。

除了使用 class 實作情境管理器,還可以透過結合contextlib模組中的@contextmanager裝飾器和函數來達到相同的效果,請看以下的例子:

from contextlib import contextmanager

@contextmanager
def open_file(filename, task):
try:
print("writing")
f = open(filename, task)
f.write("test123456")
yield f
finally:
f.close()
print("file has been closed")

with open_file("sample.txt", "w") as f:
print("file has been written.")
print(f.closed)

要理解上方的程式碼一樣從 with 語句開始閱讀

  1. 首先執行到 with 語句時,open_file function 會被調用,將 filenametask 的值傳入open_file,首先執行 print("writing") 再打開檔案寫入文字,此時碰到了 yield,function 的執行會先暫停
  2. yield 的作用有點像是 return , 執行 return 會結束 function 並且拋出 return 後的值,而 yield 不會結束 function,會先暫時將 f 拋出去
  3. 此時從 open_file function 拋出去的 f 變數會被 with 語句的 as f 所接住,接下來會執行 print("file has been written.")
  4. 離開 with 區塊後,會回到 open_file 內執行 finally 內的程式碼,也就是關閉文件與 print("file has been closed")
  5. 最後再透過 print(f.closed) 檢查檔案是否如期關閉

實際執行的結果如下,與剛才實作 class 情境管理器的結果是相同的

writing
file has been written.
file has been closed
True

總結:

在本次的教學中提到了情境管理器的介紹、優點,以及如何用 class 與 function 實作,相信讀者能夠發現使用情境管理器撰寫的程式碼有更高的閱讀性以及更好的防呆機制。 本次教學僅說明基本的情境管理器操作,未來有機會會再講解更進階的操作。

參考資料:
https://www.youtube.com/watch?v=-aKFBoZpiqA&ab_channel=CoreySchafer

--

--