Python 工匠：使用裝飾器的技巧

前言

這是 「Python 工匠」系列的第 8 篇文章。[查看系列全部文章]

裝飾器*（Decorator）* 是 Python 裏的一種特殊工具，它爲咱們提供了一種在函數外部修改函數的靈活能力。它有點像一頂畫着獨一無二 @ 符號的神奇帽子，只要將它戴在函數頭頂上，就能悄無聲息的改變函數自己的行爲。

你可能已經和裝飾器打過很多交道了。在作面向對象編程時，咱們就常常會用到 @staticmethod 和 @classmethod 兩個內置裝飾器。此外，若是你接觸過 click 模塊，就更不會對裝飾器感到陌生。click 最爲人所稱道的參數定義接口 @click.option(...) 就是利用裝飾器實現的。

除了用裝飾器，咱們也常常須要本身寫一些裝飾器。在這篇文章裏，我將從 最佳實踐 和 常見錯誤 兩個方面，來與你分享有關裝飾器的一些小知識。

最佳實踐

1. 嘗試用類來實現裝飾器

絕大多數裝飾器都是基於函數和 閉包 實現的，但這並不是製造裝飾器的惟一方式。事實上，Python 對某個對象是否能經過裝飾器（@decorator）形式使用只有一個要求：decorator 必須是一個「可被調用（callable）的對象。

# 使用 callable 能夠檢測某個對象是否「可被調用」
>>> def foo(): pass
...
>>> type(foo)
<class 'function'>
>>> callable(foo)
True
複製代碼

函數天然是「可被調用」的對象。但除了函數外，咱們也可讓任何一個類（class）變得「可被調用」（callable）。辦法很簡單，只要自定義類的 __call__ 魔法方法便可。

class Foo:
    def __call__(self):
        print("Hello, __call___")

foo = Foo()

# OUTPUT: True
print(callable(foo))
# 調用 foo 實例
# OUTPUT: Hello, __call__
foo()
複製代碼

基於這個特性，咱們能夠很方便的使用類來實現裝飾器。

下面這段代碼，會定義一個名爲 @delay(duration) 的裝飾器，使用它裝飾過的函數在每次執行前，都會等待額外的 duration 秒。同時，咱們也但願爲用戶提供無需等待立刻執行的 eager_call 接口。

import time
import functools


class DelayFunc:
    def __init__(self, duration, func):
        self.duration = duration
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f'Wait for {self.duration} seconds...')
        time.sleep(self.duration)
        return self.func(*args, **kwargs)

    def eager_call(self, *args, **kwargs):
        print('Call without delay')
        return self.func(*args, **kwargs)


def delay(duration):
    """裝飾器：推遲某個函數的執行。同時提供 .eager_call 方法當即執行 """
    # 此處爲了不定義額外函數，直接使用 functools.partial 幫助構造
    # DelayFunc 實例
    return functools.partial(DelayFunc, duration)
複製代碼

如何使用裝飾器的樣例代碼：

@delay(duration=2)
def add(a, b):
    return a + b


# 此次調用將會延遲 2 秒
add(1, 2)
# 此次調用將會當即執行
add.eager_call(1, 2)
複製代碼

@delay(duration) 就是一個基於類來實現的裝飾器。固然，若是你很是熟悉 Python 裏的函數和閉包，上面的 delay 裝飾器其實也徹底能夠只用函數來實現。因此，爲何咱們要用類來作這件事呢？

與純函數相比，我以爲使用類實現的裝飾器在特定場景下有幾個優點：

• 實現有狀態的裝飾器時，操做類屬性比操做閉包內變量更符合直覺、不易出錯
• 實現爲函數擴充接口的裝飾器時，使用類包裝函數，比直接爲函數對象追加屬性更易於維護
• 更容易實現一個同時兼容裝飾器與上下文管理器協議的對象（參考 unitest.mock.patch）

2. 使用 wrapt 模塊編寫更扁平的裝飾器

在寫裝飾器的過程當中，你有沒有碰到過什麼不爽的事情？無論你有沒有，反正我有。我常常在寫代碼的時候，被下面兩件事情搞得特別難受：

1. 實現帶參數的裝飾器時，層層嵌套的函數代碼特別難寫、難讀
2. 由於函數和類方法的不一樣，爲前者寫的裝飾器常常無法直接套用在後者上

好比，在下面的例子裏，我實現了一個生成隨機數並注入爲函數參數的裝飾器。

import random


def provide_number(min_num, max_num):
    """裝飾器：隨機生成一個在 [min_num, max_num] 範圍的整數，追加爲函數的第一個位置參數 """
    def wrapper(func):
        def decorated(*args, **kwargs):
            num = random.randint(min_num, max_num)
            # 將 num 做爲第一個參數追加後調用函數
            return func(num, *args, **kwargs)
        return decorated
    return wrapper
    


@provide_number(1, 100)
def print_random_number(num):
    print(num)

# 輸出 1-100 的隨機整數
# OUTPUT: 72
print_random_number()
複製代碼

@provide_number 裝飾器功能看上去很不錯，但它有着我在前面提到的兩個問題：**嵌套層級深、沒法在類方法上使用。**若是直接用它去裝飾類方法，會出現下面的狀況：

class Foo:
    @provide_number(1, 100)
    def print_random_number(self, num):
        print(num)

# OUTPUT: <__main__.Foo object at 0x104047278>
Foo().print_random_number()
複製代碼

Foo 類實例中的 print_random_number 方法將會輸出類實例 self ，而不是咱們指望的隨機數 num。

之因此會出現這個結果，是由於類方法*（method）和函數（function）*兩者在工做機制上有着細微不一樣。若是要修復這個問題，provider_number 裝飾器在修改類方法的位置參數時，必須聰明的跳過藏在 *args 裏面的類實例 self 變量，才能正確的將 num 做爲第一個參數注入。

這時，就應該是 wrapt 模塊閃亮登場的時候了。wrapt 模塊是一個專門幫助你編寫裝飾器的工具庫。利用它，咱們能夠很是方便的改造 provide_number 裝飾器，完美解決*「嵌套層級深」和「沒法通用」*兩個問題，

import wrapt

def provide_number(min_num, max_num):
 @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        # 參數含義：
        #
        # - wrapped：被裝飾的函數或類方法
        # - instance：
        # - 若是被裝飾者爲普通類方法，該值爲類實例
        # - 若是被裝飾者爲 classmethod 類方法，該值爲類
        # - 若是被裝飾者爲類/函數/靜態方法，該值爲 None
        #
        # - args：調用時的位置參數（注意沒有 * 符號）
        # - kwargs：調用時的關鍵字參數（注意沒有 ** 符號）
        #
        num = random.randint(min_num, max_num)
        # 無需關注 wrapped 是類方法或普通函數，直接在頭部追加參數
        args = (num,) + args
        return wrapped(*args, **kwargs)
    return wrapper
    
<... 應用裝飾器部分代碼省略 ...>
    
# OUTPUT: 48
Foo().print_random_number()
複製代碼

使用 wrapt 模塊編寫的裝飾器，相比原來擁有下面這些優點：

• 嵌套層級少：使用 @wrapt.decorator 能夠將兩層嵌套減小爲一層
• 更簡單：處理位置與關鍵字參數時，能夠忽略類實例等特殊狀況
• 更靈活：針對 instance 值進行條件判斷後，更容易讓裝飾器變得通用

常見錯誤

1. 「裝飾器」並非「裝飾器模式」

「設計模式」是一個在計算機世界裏鼎鼎大名的詞。假如你是一名 Java 程序員，而你一點設計模式都不懂，那麼我打賭你找工做的面試過程必定會度過的至關艱難。

但寫 Python 時，咱們極少談起「設計模式」。雖然 Python 也是一門支持面向對象的編程語言，但它的 鴨子類型 設計以及出色的動態特性決定了，大部分設計模式對咱們來講並非必需品。因此，不少 Python 程序員在工做很長一段時間後，可能並無真正應用過幾種設計模式。

不過 「裝飾器模式（Decorator Pattern）」 是個例外。由於 Python 的「裝飾器」和「裝飾器模式」有着如出一轍的名字，我不止一次聽到有人把它們倆當成一回事，認爲使用「裝飾器」就是在實踐「裝飾器模式」。但事實上，它們是兩個徹底不一樣的東西。

「裝飾器模式」是一個徹底基於「面向對象」衍生出的編程手法。它擁有幾個關鍵組成：一個統一的接口定義、若干個遵循該接口的類、類與類之間一層一層的包裝。最終由它們共同造成一種*「裝飾」*的效果。

而 Python 裏的「裝飾器」和「面向對象」沒有任何直接聯繫，**它徹底能夠只是發生在函數和函數間的把戲。**事實上，「裝飾器」並無提供某種沒法替代的功能，它僅僅就是一顆「語法糖」而已。下面這段使用了裝飾器的代碼：

@log_time
@cache_result
def foo(): pass
複製代碼

基本徹底等同於下面這樣：

def foo(): pass

foo = log_time(cache_result(foo))
複製代碼

裝飾器最大的功勞，在於讓咱們在某些特定場景時，能夠寫出更符合直覺、易於閱讀的代碼。它只是一顆「糖」，並非某個面向對象領域的複雜編程模式。

Hint: 在 Python 官網上有一個 實現了裝飾器模式的例子，你能夠讀讀這個例子來更好的瞭解它。

2. 記得用 functools.wraps() 裝飾內層函數

下面是一個簡單的裝飾器，專門用來打印函數調用耗時：

import time


def timer(wrapped):
    """裝飾器：記錄並打印函數耗時"""
    def decorated(*args, **kwargs):
        st = time.time()
        ret = wrapped(*args, **kwargs)
        print('execution take: {} seconds'.format(time.time() - st))
        return ret
    return decorated


@timer
def random_sleep():
    """隨機睡眠一小會"""
    time.sleep(random.random())
複製代碼

timer 裝飾器雖然沒有錯誤，可是使用它裝飾函數後，函數的原始簽名就會被破壞。也就是說你再也沒辦法正確拿到 random_sleep 函數的名稱、文檔內容了，全部簽名都會變成內層函數 decorated 的值：

print(random_sleep.__name__)
# 輸出 'decorated'
print(random_sleep.__doc__)
# 輸出 None
複製代碼

這雖然只是個小問題，但在某些時候也可能會致使難以察覺的 bug。幸運的是，標準庫 functools 爲它提供瞭解決方案，你只須要在定義裝飾器時，用另一個裝飾器再裝飾一下內層 decorated 函數就行。

聽上去有點繞，但其實就是新增一行代碼而已：

def timer(wrapped):
    # 將 wrapper 函數的真實簽名賦值到 decorated 上
 @functools.wraps(wrapped)
    def decorated(*args, **kwargs):
        # <...> 已省略
    return decorated
複製代碼

這樣處理後，timer 裝飾器就不會影響它所裝飾的函數了。

print(random_sleep.__name__)
# 輸出 'random_sleep'
print(random_sleep.__doc__)
# 輸出 '隨機睡眠一小會'
複製代碼

3. 修改外層變量時記得使用 nonlocal

裝飾器是對函數對象的一個高級應用。在編寫裝飾器的過程當中，你會常常碰到內層函數須要修改外層函數變量的狀況。就像下面這個裝飾器同樣：

import functools

def counter(func):
    """裝飾器：記錄並打印調用次數"""
    count = 0
 @functools.wraps(func)
    def decorated(*args, **kwargs):
        # 次數累加
        count += 1
        print(f"Count: {count}")
        return func(*args, **kwargs)
    return decorated

@counter
def foo():
    pass

foo()
複製代碼

爲了統計函數調用次數，咱們須要在 decorated 函數內部修改外層函數定義的 count 變量的值。可是，上面這段代碼是有問題的，在執行它時解釋器會報錯:

Traceback (most recent call last):
  File "counter.py", line 22, in <module>
    foo()
  File "counter.py", line 11, in decorated
    count += 1
UnboundLocalError: local variable 'count' referenced before assignment
複製代碼

這個錯誤是由 counter 與 decorated 函數互相嵌套的做用域引發的。

當解釋器執行到 count += 1 時，並不知道 count 是一個在外層做用域定義的變量，它把 count 當作一個局部變量，並在當前做用域內查找。最終卻沒有找到有關 count 變量的任何定義，而後拋出錯誤。

爲了解決這個問題，咱們須要經過 nonlocal 關鍵字告訴解釋器：「count 變量並不屬於當前的 local 做用域，去外面找找吧」，以前的錯誤就能夠獲得解決。

def decorated(*args, **kwargs):
    nonlocal count
    count += 1
    # <... 已省略 ...>
複製代碼

Hint：若是要了解更多有關 nonlocal 關鍵字的歷史，能夠查閱 PEP-3104

總結

在這篇文章裏，我與你分享了有關裝飾器的一些技巧與小知識。

一些要點總結：

• 一切 callable 的對象均可以被用來實現裝飾器
• 混合使用函數與類，能夠更好的實現裝飾器
• wrapt 模塊頗有用，用它能夠幫助咱們用更簡單的代碼寫出複雜裝飾器
• 「裝飾器」只是語法糖，它不是「裝飾器模式」
• 裝飾器會改變函數的原始簽名，你須要 functools.wraps
• 在內層函數修改外層函數的變量時，須要使用 nonlocal 關鍵字

看完文章的你，有沒有什麼想吐槽的？請留言或者在 項目 Github Issues 告訴我吧。

附錄

• 題圖來源: Photo by Clem Onojeghuo on Unsplash
• 更多系列文章地址：github.com/piglei/one-…

系列其餘文章：

• 全部文章索引 [Github]
• Python 工匠：編寫條件分支代碼的技巧
• Python 工匠：異常處理的三個好習慣
• Python 工匠：編寫地道循環的兩個建議