Skip to content

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

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 的語法。

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

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


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

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

References

Read More

[Python] @property 裝飾器將類別方法轉為只讀屬性

Tags:

Leave a Reply