Python 裝飾器 (Decorator)

現在的 Python 程式常常會看到以 @ 開頭的程式, 有些時候看起來很玄, 不過因為太多地方會用到這樣的語法, 因此即使想不去瞭解都不行, 以下就來認識一下所謂的『裝飾器』。

裝飾器的緣起

def works():
    total = 0
    for i in range(10000):
        total += i
    print("total:", total)

works()

在測試的過程中, 我們高度懷疑這個函式效能不彰, 因此希望能幫這個函式計時, 一般的作法可能是這樣:

import time
def timing(func):
    print("Start...")
    t1 = time.perf_counter()
    func()
    t2 = time.perf_counter()
    print("Elapsed time(secs):", t2 - t1)

def works():
    total = 0
    for i in range(10000):
        total += i
    print("total:", total)

timing(works)

我們加上一個計時函式, 但為了使用這個計時函式, 所有叫用 works 的地方都要改成 timing(works), 這樣顯然很麻煩, 除了要一一修改程式外, 如果有一天不再需要計時時, 就得一一回頭把程式改回原貌。為了避免這個問題, 我們可以這樣做:

import time
def timing(func):
    def wrapper():
        print("Start...")
        t1 = time.perf_counter()
        func()
        t2 = time.perf_counter()
        print("Elapsed time(secs):", t2 - t1)
    return wrapper

def works():
    total = 0
    for i in range(10000):
        total += i
    print("total:", total)

works = timing(works)
works()

上述程式利用一個包裝函式 wrapper 把原本的函式包裝起來, 並且在倒數第 2 列將包裝過的函式再重新命名回原來被包裝的函式名稱。如此一來, 就等於是為原始的函式加上了計時的功能, 同時又可以維持原始的函式名稱, 程式中所有叫用 works 函式的地方都不需要修改, 就會自動變成計時版的 works。往後若不需要計時的功能時, 只要刪除倒數第 2 列, 就可以讓 works 變回原始的版本。

裝飾器 (decorator) 語法

對於這樣的需求, Python 提供有一種簡便的語法蜜糖 (syntax suger) 叫做裝飾器 (decorator)

import time
def timing(func):
    def wrapper():
        print("Start...")
        t1 = time.perf_counter()
        func()
        t2 = time.perf_counter()
        print("Elapsed time(secs):", t2 - t1)
    return wrapper

@timing
def works():
    total = 0
    for i in range(10000):
        total += i
    print("total:", total)

works()

@ 開頭那一列的意思就是將 works 傳入 timing 後, 再將傳回的物件重新命名為 works, 也就等於前一個程式中的倒數第 2 列。你可以把 @timing 視為 加上 timing 功能的意思。利用這樣的語法, 完全不需要更動原本叫用 works 函式的程式, 就可以幫所有叫用 works 函式的程式計時了。

裝飾器的實際動作

裝飾器的語法實際上會轉成等效的 Python 程式, 例如:

@decorator_here
def func:
    ....

會變成等效的以下程式:

func = decorator_here(func)

其中的 decorator_here 必須是可叫用 (callable) 的物件, 並且符合必須傳入單一引數及傳回可叫用物件的規範, 而裝飾器就會把原始函式的名稱繫結到叫用此物件得到的傳回值。

這樣講有點抽象, 後續說明需要參數的裝飾器時再進一步說明。

裝飾需要參數的函式

如果要套用裝飾器的函式具有參數, 那麼可以在定義包裝函式時利用 *args**kwargs 來接收任意數量的參數, 再轉傳給被包裝的函式, 例如:

import time
def timing(func):
    def wrapper(*args, **kwargs):
        print("Start...")
        t1 = time.perf_counter()
        func(*args, **kwargs)
        t2 = time.perf_counter()
        print("Elapsed time(secs):", t2 - t1)
    return wrapper

@timing
def works(start, stop):
    total = 0
    for i in range(start, stop):
        total += i
    print("total:", total)

works(1, 10000)

這樣一來, 不管是那一個函式, 都可以套用 timing 裝飾器。如果要套用裝飾器的函式有傳回值, 也一樣在包裝函式中先接收傳回值後再轉傳出來即可, 例如:

import time
def timing(func):
    def wrapper(*args, **kwargs):
        print("Start...")
        t1 = time.perf_counter()
        res = func(*args, **kwargs)
        t2 = time.perf_counter()
        print("Elapsed time(secs):", t2 - t1)
        return res
    return wrapper

@timing
def works(start, stop):
    total = 0
    for i in range(start, stop):
        total += i
    return total

print("total:", works(1, 10000))

需要參數的裝飾器

假如我們想要在每次執行這個程式時可以計算該函式執行不同次數的時間, 就會希望可以在套用裝飾器時指定次數, 那就可以再用一層函式包裝原本的裝飾器, 例如:

import time
def timing(times):
    def outer(func):
        def wrapper():
            print("Start...")
            t1 = time.perf_counter()
            for i in range(times):
                func()
            t2 = time.perf_counter()
            print("Elapsed time(secs):", t2 - t1)
        return wrapper
    return outer

@timing(2)
def works():
    print("running...")
    total = 0
    for i in range(10000):
        total += i

print("total:", works())

這裡我們在原本的裝飾器結構上再加了一層函式, 並且在叫用裝飾器時傳遞了代表執行次數的參數, 這時會依照裝飾器的實際動作中所提到, 變成以下的等效程式:

works = timing(2)(works)

這裡 timing(2) 會傳回 outer 函式, 而 outer 函式符合必須傳入單一引數及傳回可叫用物件的規範, 因此上述裝飾器語法的等效程式就可以正確運作。

要注意的是, 在傳回 outer 時, 會將隨 timing(2) 傳入 2 所繫結的 times 名稱納入閉包 (closure), 因此後續叫用 outer 時就仍然可以透過 times 取得指定的執行次數, 據此傳回最後的 wrapper 函式。

在 Python 3.9 之前的語法中, 雖然在裝飾器語法中可以叫用函式, 但規定就只能有一層的叫用, 例如你不能這樣使用裝飾器語法:

def deco_multiple():
    def outmost():
        def outer(func):
            def wrapper():
                print('start')
                func()
                print('done')
            return wrapper
        return outer
    return outmost

@deco_multiple()()
def works():
    print('working.')

works()

雖然 deco_multiple() 會傳回 outmost, 因此 deco_multiple()() 就等於 outmost() 會傳回 outer, 而 outer 符合必須傳入單一引數及傳回可叫用物件的規範, 但是裝飾器語法的規格就是 @ 後面只能出現一次函式叫用, 因此上面的程式執行時會出現語法錯誤:

$ py .\deco_multi_test.py
  File "dec.py", line 12
    @deco_multiple()()
                    ^
SyntaxError: invalid syntax

不過這個規定在 Python 3.9 中已經鬆綁了, 現在你可以使用任何可以用在指派運算右邊的運算式作為裝飾器, 只要該運算式的結果符合裝飾器的規範即可。因此上面的例子若在 Python 3.9 是可以正確運作的:

$ py .\deco_multi_test.py
start
working.
done

裝飾器的用途

裝飾器常用來實作登記處理函式, 例如 GUI 介面的事件處理函式, 或是網頁伺服器路徑的處理函式。

使用裝飾器處理事件

底下我們舉個簡單的例子, 來模擬 GUI 介面的事件處理函式:

class EventHandler:
    def __init__(self):
        self.handlers = {}
    def on(self, ev):
        def wrapper(func):
            self.handlers[ev] = func
            return func
        return wrapper

    def dispatch(self, ev):
        if ev in self.handlers:
            self.handlers[ev]()

這個 EventHandler 類別可以登記不同名稱事件的對應處理函式, 其中的 on 函式就是負責登記處理函式, 我們將它實作成裝飾器的形式;dispatch 函式則是會依據發生的事件, 叫用已註冊處理該事件的函式。

接著, 就可以產生 EventHandler 類別的物件, 並且使用裝飾器註冊各種事件的處理函式:

e = EventHandler()

@e.on('click')
def onClick():
    print('clicked')

@e.on('double_click')
def onDoubleClick():
    print('double clicked')

之後, 就可以叫用 dispatch 處理個別事件:

e.dispatch('click')
e.dispatch('double_click')

若這個程式名稱為 deco_event.py, 執行結果如下:

$ python .\deco_event.py
clicked
double clicked

利用裝飾器的語法, 可以很清楚的表達 onClick 和 onDoubleClick 是專門為了 e 這個物件所撰寫的事件處理函式, 也能明確的看到個別處理的事件名稱。

裝飾器語法與傳統語法並存

如果你擔心上述事件處理的方式對於不熟悉裝飾器語法的人來說有點難懂, 也可以實作成以下的變化方式:

class EventHandler:
    def __init__(self):
        self.handlers = {}
    def on(self, ev, func = None):
        if func:
            self.handlers[ev] = func
        else:
            def wrapper(func):
                self.handlers[ev] = func
                return func
            return wrapper

    def dispatch(self, ev):
        if ev in self.handlers:
            self.handlers[ev]()

e = EventHandler()

@e.on('click')
def onClick():
    print('clicked')

def onDoubleClick():
    print('double clicked')
e.on('double_click', onDoubleClick)

e.dispatch('click')
e.dispatch('double_click')

使用裝飾器語法時, 因為叫用 on() 時沒有傳遞函式給 func 參數, 因此會執行 else 的部分, 也就是原本裝飾器的實作內容;但若直接叫用 on() 時就可以同時傳遞事件處理函式, 兩種語法並存, 都可以達到相同的效果。

裝飾器不一定要改變原始函式的功能

你也可以發現到, 在這個例子中, 雖然裝飾器幫函式加上了可以處理特定事件的功能, 但是我們並沒有修改函式的內容, 實際上要如何運用裝飾器, 並沒有什麼限制, 完全看你自己的創意。

13