一種自頂而下的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. 完工！

我分享編程感悟與學習資料的公衆號，敬請關注：程序員甜甜圈