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

Last updated on  Mar 10, 2023  in  Python 程式設計 - 中階  by  Amo Chen  ‐ 6 min read

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/

對抗久坐職業傷害

研究指出每天增加 2 小時坐著的時間,會增加大腸癌、心臟疾病、肺癌的風險,也造成肩頸、腰背疼痛等常見問題。

然而對抗這些問題,卻只需要工作時定期休息跟伸展身體即可!

你想輕鬆改變現狀嗎?試試看我們的 PomodoRoll 番茄鐘吧! PomodoRoll 番茄鐘會根據你所設定的專注時間,定期建議你 1 項辦公族適用的伸展運動,幫助你打敗久坐所帶來的傷害!

贊助我們的創作

看完這篇文章了嗎? 休息一下,喝杯咖啡吧!

如果你覺得 MyApollo 有讓你獲得實用的資訊,希望能看到更多的技術分享,邀請你贊助我們一杯咖啡,讓我們有更多的動力與精力繼續提供高品質的文章,感謝你的支持!