看了肯定會的 Python 裝飾子(decorator)教學

Python 的 decorator (或稱裝飾子)是一個非常有用的功能，它的重要程度可以說是沒用過或不會用 decorator 的人就等於沒學過 Python, 甚至在一些常見的框架(Framework)，例如 Flask, FastAPI, Django 都提供各種方便的 decorator 供大家使用。

這麼重要的東西，肯定是闖江湖都會用到的金創藥啊！

但如果你剛接觸 Python 就看到類似以下裝飾子的範例，絕對會腦袋打結，為什麼函式前面還要加個 @debug 還有 @cache, 而且還很神奇能運作：

@debug
@cache
def sum(a, b):
    return a + b

本文就教大家如何理解 Python 的 decorator!

本文環境 #

• Python 3

開始之前 #

剛學程式的人，如果想要除錯 1 個函式(function)，大多都會選擇在函式內加上 print 吧，通常都會將傳入值與回傳值都列印出來，例如：

def sum(a, b):
    print('a =', a)
    print('b =', b)
    print('a + b =', a + b)
    return a + b

然而，這樣子只能針對 sum 這個函式除錯，而且每次都要做侵入式的修改，相當不便⋯⋯。

我們都知道 Python 的傳入與回傳值可以是函式，那麼也許可以做一個函式叫 debug, 然後接受任何函式傳入，並且把我們的 debug 功能加在裡面，這樣就能夠針對任何函式進行除錯，例如：

def sum(a, b):
    return a + b

def debug(func):
    print('接到 func', func.__name__)
    # 我想在這裡 debug
    return func

debug_sum = debug(sum)
debug_sum(1, 2)

p.s. 每個 function 都有 1 個屬性 __name__ 可以取得函式名稱

執行結果：

接到 func sum

可是上述範例 debug 函式執行過程並沒有辦法呼叫傳入的 func ，因為我們無法攔截到未來使用者怎麼呼叫傳進去的 func (也就是 sum 函式），也就是 debug_sum(1, 2) 沒攔截到傳進去的 1, 2，所以根本無法完成除錯的目標，但好消息是我們可以攔截到 sum 被傳進去了。

怎麼解決攔截傳入參數的部分？ #

很簡單！做個中間人！

我們可以把傳進 debug 函式的 func 再用 def 包裝在 1 個新的函式裡面，這個新的函式要負責接受任何參數，然後轉一手再代入先前傳進來的函式，最後我們再將新包裝過的函式回傳回去，於是程式碼會變成這樣：

def sum(a, b):
    return a + b

def debug(func):
    print('接到 func', func.__name__)
    def wrapper(*args, **kwargs):
        print('幫忙代入 args', args)
        print('幫忙代入 kwargs', kwargs)
        func(*args, **kwargs)
    return wrapper

debug_sum = debug(sum)
debug_sum(1, 2)
print('debug_sum', debug_sum)

執行結果：

接到 func sum
幫忙代入 args (1, 2)
幫忙代入 kwargs {}
debug_sum <function debug.<locals>.wrapper at 0x107966430>

耶！從上面執行結果看到我們透過再包裝 1 層函式，就能攔截傳入參數的部分。從執行結果也可以看到呼叫完 debug 函式，所得到的 debug_sum 其實就是那個新包裝好的函式，因此上述範例的呼叫 debug_sum 的方式可以進一步濃縮為 1 行：

debug(sum)(1, 2)

不過上述範例還是沒有列印 sum 執行結果的部分，所以可以再進一步改成：

def sum(a, b):
    return a + b

def debug(func):
    print('接到 func', func.__name__)
    def wrapper(*args, **kwargs):
        print('幫忙代入 args', args)
        print('幫忙代入 kwargs', kwargs)
        result = func(*args, **kwargs)
        print(func.__name__, '執行結果', result)
    return wrapper

debug_sum = debug(sum)
debug_sum(1, 2)

執行結果：

接到 func sum
幫忙代入 args (1, 2)
幫忙代入 kwargs {}
sum 執行結果 3

耶！完美！

完成 debug 函式之後，你就可以任意 debug 函式，而且不需要做到任何侵入性的修改，例如 debug print 函式：

def debug(func):
    print('接到 func', func.__name__)
    def wrapper(*args, **kwargs):
        print('幫忙代入 args', args)
        print('幫忙代入 kwargs', kwargs)
        result = func(*args, **kwargs)
        print(func.__name__, '執行結果', result)
    return wrapper

debug_print = debug(print)
debug_print(1, 2, '3')

看到這邊，大家應該就會了解 Python decorator 在做什麼 — 接收 1 個函式後加工並包裝成 1 個新的函式！

@ 語法糖(Syntactic sugar) #

也由於這種 decorator 的模式很好用，所以 Python 提供 1 種特殊的語法 @ 可以讓我們將 debug_sum = debug(sum) 簡化為：

@debug
def sum(a, b):
    return a + b

這種語法就被稱為是 1 種語法糖，太好吃了，而且會上癮⋯⋯。

所以這種接受接一個函式作為參數，然後返回一個新的函式的函式都可以在前面加上 @ 作為裝飾子使用，例如下列範例的 timeit （測量函式執行時間）：

import time


def timeit(func):
    def wrapper():
        s = time.time()
        func()
        print(func.__name__, 'total time', time.time() - s)
    return wrapper


@timeit
def sleep_10s():
    time.sleep(10)


sleep_10s()

執行結果：

sleep_10s total time 10.000677108764648

多個裝飾子的執行順序 #

由於裝飾子非常好用，所以你很有可能會看到類似的程式碼片段存在專案之中：

@timeit
@api
@auth_required
@cache
def get_profile():
    ...(略)

這時候頭就大了，到底要怎麼理解多個裝飾子的執行順序？到底是 @timeit 會先執行？還是 @cache 會先執行呢？

為了解答這個問題，我們可以將 decorator 的數量簡化至 2 個，例如：

def deco1(func):
    def wrapper():
        print('deco1')
        func()
        print('deco1 end')
    return wrapper


def deco2(func):
    def wrapper():
        print('deco2')
        func()
        print('deco2 end')
    return wrapper


@deco1
@deco2
def main():
    print('main')


main()

上述執行結果如下：

deco1
deco2
main
deco2 end
deco1 end

其實從執行結果可以看出，最外層的 decorator 會最先執行，最晚結束。

但其實也不難理解，因為 decorator 是一層包一層的形式，所以只要把一層包一層的圖畫出來，再畫一條直線從上往下貫穿，我們就可以理解其執行與結束順序了：

functools.wraps #

先前章節提到 decorator 會接受接一個函式作為參數，然後返回一個新的函式，這其實會有 1 個小小問題產生，就是被包裝的函式的名字與 doc string 都會消失，例如：

import time


def timeit(func):
    def wrapper():
        s = time.time()
        func()
        print(func.__name__, 'total time', time.time() - s)
    return wrapper


@timeit
def sleep_10s():
    """sleep 10s"""
    time.sleep(10)


print('func', sleep_10s.__name__)
print('doc', sleep_10s.__doc__)

執行結果：

func wrapper
doc string None

從上述執行結果可以看到 sleep_10s 的名字與 doc string 都不是我們所預期的樣子，原因在於 decorator 其實回傳新的函式，所以這些非預期的值都數於新回傳的函式，這樣就對開發協作者比較不友善了⋯⋯。

如果要修好這個問題，可以用 functools.wrap 再包裝一次回傳的函式：

import time

from functools import wraps


def timeit(func):
    @wraps(func)
    def wrapper():
        s = time.time()
        func()
        print(func.__name__, 'total time', time.time() - s)
    return wrapper


@timeit
def sleep_10s():
    """sleep 10s"""
    time.sleep(10)


print('func', sleep_10s.__name__)
print('doc', sleep_10s.__doc__)

執行結果如下，可以看到 sleep_10s 的函式名稱與 doc string 又恢復正常了：

func sleep_10s
doc string sleep 10s

實際上 functool.wraps 可以視情況自行決定要不要加，如果是一定要保留 doc string 或原本的函式名稱的話，就可以用 functool.wraps, 否則其實不加也不影響日常使用。

類別裝飾子(Class-based decorators) #

其實，不只有函式(function)能夠當裝飾子，類別(class)也能改裝成裝飾子。

類別寫成的裝飾子就稱為 class-based decorator 。

實作 classed-based decorator 的方法雖然不如以 function 方式實作直覺，但也很簡單，只要實作 __init__() 方法接受傳入函式，並且實作 __call__ 方法呼叫被傳入的函式即可，例如：

class ClsDeco:
    def __init__(self, func):
        self.func = func

    def __call__(self):
        print("before calling", self.func.__name__)
        self.func()
        print("after calling", self.func.__name__)

@ClsDeco
def say_hello():
    print("Hello!")

say_hello()

執行結果：

before calling say_hello
Hello!
after calling say_hello

前述範例其實等同於下列形式， 也就是 ClsDeco(say_hello)() 的部分：

class ClsDeco:
    def __init__(self, func):
        self.func = func

    def __call__(self):
        print("before calling", self.func.__name__)
        self.func()
        print("after calling", self.func.__name__)


def say_hello():
    print("Hello!")


ClsDeco(say_hello)()

這就是關於 classed-based decorator 的簡單介紹。

更複雜的 decorator - 可設定參數的裝飾子 #

學會 decorator 與 class-based decorator 之後，還可以進一步製造出更複雜的裝飾子，譬如接受參數設定的裝飾子，例如下列範例 @retry 裝飾子接受參數 max=3 ，改變裝飾子的行為：

@retry(max=3)
def get_stock_price():
    pass

要怎麼解讀 @retry(max=3) 呢？其實就是 1 個函數回傳另 1 個裝飾子：

r = retry(max=3)

@r
def get_stock_price():
    pass

實作上就類似下列的程式碼：

def retry(max=1):
    class Wrapper:
        def __init__(self, func):
            self.func = func

        def __call__(self):
            retried = 0
            while retried < max:
                 try:
                     self.func()
                 except Exception:
                     retried += 1
                     print('Failed. Going to try again (', retried, ')')
                 else:
                     break
    return Wrapper


@retry(max=3)
def get_stock_price():
    raise ValueError


get_stock_price()

執行結果如下：

Failed. Going to try again ( 1 )
Failed. Going to try again ( 2 )
Failed. Going to try again ( 3 )

從上述結果可以看到，我們藉著 1 個函數回傳另 1 個裝飾子的做法，成功讓裝飾子具有可設定的特性，不過相對也讓程式變了複雜一點。

以上就是關於更複雜的 decorator 的介紹。

總結 #

Decorator 是一個非常實用的模式/功能，它可以讓我們輕鬆地在既有基礎上疊加額外的功能，除了使程式碼更加簡潔、易讀之外，還可以增加複用性。

不過由於 Decorator 可以不斷疊加的特性，甚至是可以多重包裝一個函式，有時候會適得其反，造成程式碼閱讀困難，使用上還是建議盡量保持單純為佳。

如果你還沒有在 Python 專案中使用 decorator, 現在是時候了！

以上！ Happy Coding!

References #

https://peps.python.org/pep-0318/