一種自頂而下的Python裝飾器設計方法

追夢人物發表於2019-02-25
Python

本文首發於我的個人部落格，更多 Python、django 和 Vue 的開發教程，請訪問 追夢人物的部落格。

裝飾器是 Python 的一種重要程式設計實踐，然而如果沒有掌握其原理和適當的方法，寫 Python 裝飾器時就可能遇到各種困難。猶記得當年校招時應聘今日頭條 Python 開發崗位，因一道 Python 裝飾器的設計問題而止於終面，非常遺憾。隨著程式設計技術的提升以及對 Python 裝飾器更加深入的理解，我逐漸總結出一套自頂而下的裝飾器設計方法，這個方法能夠指導我們輕鬆寫出各種型別的裝飾器，再也不用像以前那樣死記硬背裝飾器的模板程式碼。

Python裝飾器原理

下面是 Python 裝飾器的常規寫法：

@decorator
def func(*args, **kwargs):
    do_something()
複製程式碼

這種寫法只是一種語法糖，使得程式碼看起來更加簡潔而已，在 Python 直譯器內部，函式 func 的呼叫被轉換為下面的方式：

>>> func(a, b, c='value')
# 等價於
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
複製程式碼

可見，裝飾器 decorator 是一個函式（當然也可以是一個類），它接收被裝飾的函式 func 作為唯一的引數，然後返回一個 callable（可呼叫物件），對被裝飾函式 func 的呼叫實際上是對返回的 callable 物件的呼叫。

自頂而下設計裝飾器

從原理分析可見，如果我們要設計一個裝飾器，將原始的函式（或類）裝飾成一個功能更加強大的函式（或類），那麼我們要做的就是要寫一個函式（或類），其被呼叫後返回我們需要的那個功能更加強大的函式（或類）。

簡單裝飾器

簡單的裝飾器函式就像上面介紹的那樣，不帶任何引數。假設我們要設計一個裝飾器函式，其功能是能使得被裝飾的函式呼叫結束後，列印出函式執行時間，我們來看看使用自頂而下的方法來設計這個裝飾器該怎麼做。

所謂“頂”，就是先不關注實現細節，而是做好整體設計和分解函式呼叫過程。我們把裝飾器命名為 timethis，其使用方法像下面這樣：

@timethis
def fun(*args, **kwargs):
    pass
複製程式碼

分解對被裝飾函式 fun 的呼叫過程：

>>> func(*args, **kwargs)
# 等價於
>>> decorated_func = timethis(func)
>>> decorated_func(a, b, c='value')
複製程式碼

由此可見，我們的裝飾器 timethis 應該接收被裝飾的函式作為唯一引數，返回一個函式物件，根據慣例，返回的函式命名為 wrapper，因此可以寫出 timethis 裝飾器的模板程式碼：

def timethis(func):
    def wrapper(*args, **kwargs):
        pass
    
    return wrapper
複製程式碼

裝飾器的框架搭好了，接下來就是“下”，豐富函式邏輯。

對被裝飾的函式呼叫等價於對 wrapper 函式的呼叫，為了使 wrapper 呼叫返回和被裝飾函式呼叫一樣的結果，我們可以在 wrapper 中呼叫原函式並返回其呼叫結果：

def timethis(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
    	return result
    return wrapper
複製程式碼

可以隨意豐富 wrapper 函式的邏輯，我們的需求是列印 func 的呼叫時間，只需在 func 呼叫前後計時即可：

import time

def timethis(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end-start)
    	return result
    return wrapper
複製程式碼

由此，一個可以列印函式呼叫時間的裝飾器就完成了，來看看使用效果：

@timethis
def fibonacci(n):
    """
    求斐波拉契數列第 n 項的值
    """
    a = b = 1
    while n > 2:
        a, b = b, a + b
        n -= 1
    return b

>>> fibonacci(10000)
fibonacci 0.004000663757324219
...結果太大省略
複製程式碼

基本上看上去沒有問題了，不過由於函式被裝飾了，因此被裝飾函式的基本資訊變成了裝飾器返回的 wrapper 函式的資訊：

>>> fibonacci.__name__
wrapper
>>> fibonacci.__doc__
None
複製程式碼

注意這裡 fibonacci.__name__ 等價於 timethis(fibonacci).__name__，所以返回值為 wrapper。

修正方法也很簡單，需要使用標準庫中提供的一個 wraps 裝飾器，將被裝飾函式的資訊複製給 wrapper 函式：

from functools import wraps
import time

def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end - start)
        return result

    return wrapper
複製程式碼

至此，一個完整的，不帶引數的裝飾器便寫好了。

帶引數的裝飾器

上面設計的裝飾器比較簡單，不帶任何引數。我們也會經常看到帶引數的裝飾器，其使用方法大概如下：

@logged('debug', name='example', message='message')
def fun(*args, **kwargs):
    pass
複製程式碼

分解對被裝飾函式 fun 的呼叫過程：

>>> func(a, b, c='value')
# 等價於
>>> decorator = logged('debug', name='example', message='message')
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
複製程式碼

由此可見，logged 是一個函式，它返回一個裝飾器，這個返回的裝飾器再去裝飾 func 函式，因此 logged 的模板程式碼應該像這樣：

def logged(level, name=None, message=None):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            pass
        return wrapper
    return decorator
複製程式碼

wrapper 是最終被呼叫的函式，我們可以隨意豐富完善 decorator 和 wrapper 的邏輯。假設我們的需求是被裝飾函式 func 被呼叫前列印一行 log 日誌，程式碼如下：

from functools import wraps

def logged(level, name=None, message=None):
    def decorator(func):
        logname = name if name else func.__module__
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            print(level, logname, logmsg, sep=' - ')
            return func(*args, **kwargs)
        return wrapper
    return decorator
複製程式碼

多功能裝飾器

有時候，我們也會看到同一個裝飾器有兩種使用方法，可以像簡單裝飾器一樣使用，也可以傳遞引數。例如：

@logged
def func(*args, **kwargs):
    pass

@logged(level='debug', name='example', message='message')
def fun(*args, **kwargs):
    pass
複製程式碼

根據前面的分析，不帶引數的裝飾器和帶引數的裝飾器定義是不同的。不帶引數的裝飾器返回的是被裝飾後的函式，帶引數的裝飾器返回的是一個不帶引數的裝飾器，然後這個返回的不帶引數的裝飾器再返回被裝飾後的函式。那麼怎麼統一呢？先來分析一下兩種裝飾器用法的呼叫過程。

# 使用 @logged 直接裝飾
>>> func(a, b, c='value')
# 等價於
>>> decorated_func = logged(func)
>>> decorated_func(a, b, c='value')

# 使用 @logged(level='debug', name='example', message='message') 裝飾
>>> func(a, b, c='value')
# 等價於
>>> decorator = logged(level='debug', name='example', message='message')
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
複製程式碼

可以看到，第二種裝飾器比第一種裝飾器多了一步，就是呼叫裝飾器函式再返回一個裝飾器，這個返回的裝飾器和不帶引數的裝飾器是一樣的：接收被裝飾的函式作為唯一引數。唯一的區別是返回的裝飾器攜帶固定引數，固定函式引數正是 partial 函式的使用場景，因此我們可以定義如下的裝飾器：

from functools import wraps, partial

def logged(func=None, *, level='debug', name=None, message=None):
    if func is None:
        return partial(logged, level=level, name=name, message=message)

    logname = name if name else func.__module__
    logmsg = message if message else func.__name__

    @wraps(func)
    def wrapper(*args, **kwargs):
        print(level, logname, logmsg, sep=' - ')
        return func(*args, **kwargs)
    return wrapper
複製程式碼

實現的關鍵在於，若這個裝飾器以帶引數的形式使用，這第一個引數 func 的值為 None，此時我們使用 partial 返回了一個其它引數固定的裝飾器，這個裝飾器與不帶引數的簡裝飾器一樣，接收被裝飾的函式物件作為唯一引數，然後返回被裝飾後的函式物件。

裝飾類

由於類的例項化和函式呼叫非常類似，因此裝飾器函式也可以用於裝飾類，只是此時裝飾器函式的第一個引數不再是函式，而是類。基於自頂而下的設計方法，設計一個用於裝飾類的裝飾器函式就是輕而易舉的事情，這裡不再給出示例。

練習

最後，以當時今日頭條的面試題作為一個練習。現在看來這道題只是一個簡單的裝飾器設計需求，只怪自己學藝不精，後悔沒有早點掌握裝飾器的設計方法。

題目：

設計一個裝飾器函式 retry，當被裝飾的函式呼叫丟擲指定的異常時，函式會被重新呼叫，直到達到指定的最大呼叫次數才重新丟擲指定的異常。裝飾器的使用示例如下：

@retry(times=10, traced_exceptions=ValueError, reraised_exception=CustomException)
def str2int(s):
    pass
複製程式碼

times 為函式被重新呼叫的最大嘗試次數。

traced_exceptions 為監控的異常，可以為 None（預設）、異常類、或者一個異常類的列表。如果為 None，則監控所有的異常；如果指定了異常類，則若函式呼叫丟擲指定的異常時，重新呼叫函式，直至成功返回結果或者達到最大嘗試次數，此時重新丟擲原異常（reraised_exception 的值為 None），或者丟擲由 reraised_exception 指定的異常。

參考程式碼

要注意實現方式不止一種，以下是我的實現版本：

from functools import wraps

def retry(times, traced_exceptions=None, reraise_exception=None):
    def decorator(func):
        
        @wraps(func)
        def wrapper(*args, **kwargs):
            n = times
            trace_all = traced_exceptions is None
            trace_specified = traced_exceptions is not None
            while True:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    traced = trace_specified and isinstance(e, traced_exceptions)
                    reach_limit = n == 0

                    if not (trace_all or traced) or reach_limit:
                        if reraise_exception is not None:
                            raise reraise_exception
                        raise
                    n -= 1
        return wrapper
    return decorator
複製程式碼

總結

總結一下，自定而下設計裝飾器分以下幾個步驟

  1. 確定你的裝飾器該如何使用，帶引數或者不帶引數，還是都可以。
  2. 將 @ 語法糖分解為裝飾器的實際呼叫過程。
  3. 根據裝飾的呼叫過程，寫出對應的模板程式碼。
  4. 根據需求編寫裝飾器函式和裝飾後函式的邏輯。
  5. 完工！

我分享程式設計感悟與學習資料的公眾號，敬請關注：程式設計師甜甜圈

相關文章