Skip to content

[Python] @contextmanager 裝飾器的使用方式

Last Updated on 2024-09-11 by Clay

Python 中的上下文管理器裝飾器 @contextmanagercontextlib 模組中的一個裝飾器,可以讓開發人員簡便地創造屬於我們的上下文管理器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 的語法。

以下幾個場景非常適合上下文管理器:

  1. 釋放資源:在處理過程後一定要釋放出資源的邏輯,適合使用上下文管理器的自動退出
  2. 臨時改變變數值:假設有一應用場景,需要操作時改變變數值,操作結束後恢復原狀,則可以使用上下文管理器(我最近遇到的一個這樣的應用場景是更新資料庫,更新過程中我的系統得處於 maintain 的狀態,可是有時使用者上傳的資料格式是錯誤的... 但無論如何,我終究還是得恢復 online 的狀態
  3. 鎖(lock)的存取與釋放


總結以上幾點,我們可以理解 @contextmanager 的使用方法是:

  • @contextmanager 用來創建簡化的上下文管理器,通過 yield 來分別定義進入和退出上下文時要做的事。
  • 進入上下文時,執行 yield 之前的代碼;退出上下文時,執行 yield 後的代碼,通常是清理工作。
  • 搭配 tryfinally 能確保即使發生異常也能正確釋放資源。

References


Read More

Tags:

Leave a Reply