Python之裝飾器

本節內容

---

• 必要知識回顧
• 情景模擬
• 裝飾器的概念及實現原理
• 回馬槍（帶參數的裝飾器）

一、 必要知識回顧

---

在開始說裝飾器之前，需要大家熟悉之前說過的相關知識點：

• 函數即“變量”： 函數名就是一個變量名，它的值就是其對應的函數體；函數體也可以賦值給其它變量，通過這個變量也能調用函數；
• 嵌套函數： 函數內部可以嵌套定義（一層或多層）函數，內部函數可以在函數體內部調用，也可以當做返回值返回；
• 閉包： 在一個嵌套函數中，內部函數可以調用外部非全局變量並且不受外部函數生命周期的影響；
• 高階函數： 函數的參數可以是函數；

簡單來講，裝飾器就是對這些內容的整合和經典應用。如果不了解這些內容，可以查看 這篇文章。

二、情景模擬

---

我們將通過對一個功能需求的分析和解決過程來探究一下“裝飾器是什麼”以及“裝飾器的一些特性”。

1. 場景說明

假設我現在已經定義了一些函數，並且這些函數都已經被線上業務廣泛應用。

import time

def func1():
    time.sleep(1)
    print('func1')
    return 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

...

2. 功能需求

現在線上某個業務響應時間過長，需要在不影響線上服務的情況下分別統計這些函數的運行時間來定位故障。

需求解讀：

既然不能影響線上業務，那麼必然是不能要求函數調用方去更改代碼的，這當然包括調用方式。

實現方式：

• 為每個函數單獨添加重復的代碼
• 把統計運行時間的的代碼封裝成一個可接收函數作為參數的高階函數
• 使用嵌套函數改進上面的高階函數

3. 實現方式及改進過程

初步實現：分別為各個函數添加運行時間統計功能

import time
import sys

def func1():
    start_time = time.time()
    time.sleep(1)
    print('func1')
    end_time = time.time()
    func_name = sys._getframe().f_code.co_name
    print('%s run time is: %s' % (func_name, (end_time - start_time)))
    return 'func1'

def func2():
    start_time = time.time()
    time.sleep(2)
    print('func2')
    end_time = time.time()
    func_name = sys._getframe().f_code.co_name
    print('%s run time is: %s' % (func_name, (end_time - start_time)))
    return 'func2'
...

函數調用：

func1()
func2()

輸出結果：

func1
func1 run time is: 1.000861644744873
func2
func2 run time is: 2.0005600452423096

存在的問題：

功能是實現了，但是存在以下幾個問題：

• 如果涉及的函數太多，那麼需要做大量的重復性工作；
• 如果這些函數分散在不同的模塊，且有不同的人維護，那麼需要進行協商溝通來保證代碼功能一致性；
• 故障解決後，還需要一個一個的去刪除或注釋調試代碼，又是一次大量的重復性工作；
• 如果其他業務也遇到相同問題，需要再次在多個地方復寫這些代碼。

改進思路：

要避免重復勞動，提高代碼重用性，大家很自然就會想到把這個功能封裝成一個函數。

改進1：自定義一個（高階）函數來提高統計代碼復用性

import time

def func1():
    time.sleep(1)
    print('func1')
    return 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

def print_run_time(f):
    time_start = time.time()
    ret = f()
    time_end = time.time()
    func_name = f.__name__
    print('%s run time is: %s' % (func_name, (time_end - time_start)))
    return ret

函數調用：

print_run_time(func1)
print_run_time(func2)

輸出結果：

func1
func1 run time is: 1.0003290176391602
func2
func2 run time is: 2.0004265308380127

存在的問題：

統計代碼的重復性工作解決了，但是新的問題出現了：此時只能通過print_run_time(f)函數去調用原來的函數了（func1, func2, ...），這顯然已經改變了函數的調用方式，因此是不合理的。

看上去越改越差勁了，我們繼續分析下，會峰回路轉的。

改進思路：

如果函數調用方式不能修改，那麼我們只能給原來的函數名重新賦值一個新的函數體了，這個新的函數體就應該是添加完統計功能之後的函數體。我們看看下面這兩種實現行不行：

func1 = print_run_time
func2 = print_run_time

print_run_time(f)是有參數的，而原函數func1和func2是沒參數的，調用方式還是發生改變了，因此這種方式不可行。

func1 = print_run_time(func1)
func2 = print_run_time(func2)

上面這種方式顯然更不行了，因為print_run_time(f)返回的是原函數的返回值，而這個返回值不是一個函數，這將導致被重新賦值後的func1和func2無法被調用。

那麼我們是否可以把print_run_time(f)函數的函數體定義為一個內部的嵌套函數，然後將這個內部的嵌套函數作為print_run_time(f)函數的返回值呢？這貌似是說的通的，看下面的實現。

改進2：使用嵌套函數和閉包改進上面定義的高階函數

import time

def func1():
    time.sleep(1)
    print('func1')
    retrun 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

def print_run_time(f):
    def inner():
        time_start = time.time()
        ret = f()
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner
    
func1 = print_run_time(func1)
func2 = print_run_time(func2)

函數調用：

func1()
func2()

輸出結果：

func1
func1 run time is: 1.0003256797790527
func2
func2 run time is: 2.000091791152954

Cool! We got it! 我們現在所需要的做的是在所有需要統計運行時長的函數定義之後的任意地方執行一下下面這條語句就可以了：

funcN = print_run_time(funcN)

存在的問題：

我們貌似忽略了一個問題，如果原函數有參數怎麼辦？

改進思路：

是的，我們定義print_run_time(f)函數的內部函數inner()時，為它定義相應的參數就可以了。由於每個函數的參數數量是不同的，因此inner函��的參數應該是可變(長)參數。

改進3：支持原函數傳遞參數

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')
    
def print_run_time(f):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner

func1 = print_run_time(func1)
func2 = print_run_time(func2)

函數調用：

ret1 = func1('decorator test!')
print(ret1)
ret2 = func2()
print(ret2)

輸出結果：

decorator test!
func1
func1 run time is: 1.0001435279846191
func1 return
func2
func2 run time is: 2.0005276203155518
None

到目前為止：統計函數運行時間的功能實現了，函數原來的調用方式沒有發生改變，原函數的定義也沒有發生改變。其實這就是是裝飾器的雛形。

三. 裝飾器的概念及原理

---

1. 什麼是裝飾器？

裝飾器，是一種“語法糖”，其本質上就是個函數。

2. 裝飾器的作用

它是一個裝飾其他函數的函數，用來為其他函數添加一些額外的功能。

3. 裝飾器原則

裝飾器對被裝飾的函數應該是完全透明的，即

• 不能修改被裝飾的函數的源代碼
• 不能修改被裝飾的函數的調用方式

4. 什麼樣的函數才是裝飾器？

高階函數 + 嵌套函數 => 裝飾器

這裡的高階函數需要同時滿足以下兩個條件：

• 接收函數名作為參數 -- 可以實現在不修改被裝飾函數源代碼的情況下為其添加新的功能
• 返回內部嵌套函數的函數名 -- 可以實現不用修改函數的調用方式

5. 裝飾器實現實例

再來看下上面寫的實現代碼：

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')
    
def print_run_time(f):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner

對於這段代碼來講，print_run_time(f)就已經是一個裝飾器函數。為了不改變原函數的調用方式，我們需要把print_run_time(f)函數的返回值重新賦值給原來的函數名：

func1 = print_run_time(func1)
func2 = print_run_time(func2)

但是我們上面說過，裝飾器除了是一個函數之外，還是一個“語法糖”。“語法糖”應該是可以簡化某些操作的，事實上確實是這樣。上面的過程其實可以這樣來寫：

def print_run_time(f):
    def warpper(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return wrapper

@print_run_time
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

@print_run_time
def func2():
    time.sleep(2)
    print('func2')

是的，就是這麼簡單。

@print_run_time
def func1(string):
    ...

就相當於：

func1 = print_run_time(func1)

不要問為什麼，因為Python解釋器就是這樣執行的。但是需要注意，此時print_run_time(f)函數必須定義在被修改的函數定義之前，這個很容易理解，只要捋一下代碼執行過程就明白了。另外，我們把print_run_time的內部函數名改成了wrapper，這個是裝飾器函數的慣用名稱（當然，也可以繼續使用inner或使用任意名稱）。

那麼，現在我們可以像原來那樣調用函數了（上面的修改這對函數的調用方是完全無感知的）：

func1('Decorator Test...')
func2()

輸出結果：

Decorator Test...
func1
func1 run time is: 1.0006635189056396
func2
func2 run time is: 2.0003299713134766

四、 回馬槍

---

上面的“情景模擬”部分對於print_run_time函數的改進其實還沒有完成。出於內容銜接和便於對概念理解的目的，才將這個改進放到了這裡。現在我們返回來，看下這個print_run_time(f)函數還有什麼不足。

1. 一個新需求

現在要求在完成統計函數運行時間的基礎上，如果函數運行時間超過指定的秒數則打印提示信息，且該秒數允許自定義。

存在的問題：

print_run_time(f)作為一個函數，目前對可而接收的參數限制太大--只能接受一個函數作為參數，要接收一個時間參數就必須為該函數定義新的參數。

解決思路：
這個看起來很容易解決，只要給print_run_time()函數定義一個用於接收超時時間的參數就可以了。

2. 代碼初步實現

def print_run_time(f, timeout=1):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        run_time = time_end - time_start
        print('%s run time is: %s' % (func_name, run_time))
        if run_time > timeout:
            print('PROBLEM'.rjust(30, '>'))
        return ret
    return inner

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')

那麼此時，我們可以這樣來重新給func1和func2函數賦值並指定超時時間：

func1 = print_run_time(func1, timeout=2)
func2 = print_run_time(func2, timeout=2)

函數調用：

func1('Decorator Test...')
func2()

輸出結果：

Decorator Test...
func1
func1 run time is: 1.0006680488586426
func2
func2 run time is: 2.0003504753112793
>>>>>>>>>>>>>>>>>>>>>>>PROBLEM

可以看到，我們設置的超時時間為2秒，func1函數的運行時間為1秒多，小於2秒，因此沒有輸出錯誤信息；而func2函數的運行時間為2秒多，大於2秒，因此輸出了一個PROBLEM錯誤信息。Great! It works!

3. 代碼改進

來看下這時候“裝飾器語法糖”能正常工作嗎？

因為print_run_time(f, timeout=1)的兩個參數都是位置參數，要指定timeout的值就要先指定f的值，因此寫法是這樣的：

@print_run_time(func1, timeout=2)
def func1(string):
    ...

很抱歉，這種方式行不通，會報NameError: name 'func1' is not defined。

結合這個例子和前面可以正常執行的代碼，我們可以得出一個結論：裝飾器函數只能有一個接受函數的參數；或者說以“裝飾器語法糖”格式使用的函數，必須本身是一個裝飾器函數或者其返回值是一個裝飾器函數。

解決思路：

既然f和timeout不能同時存在，但是又必須都存在，那我們就只能在print_run_time函數中定義一個內部函數來把它們分開了。也就是說print_run_time函數需要返回一個定義在它內部的函數，且這個內部函數需要滿足一個正常的裝飾器函數。那麼參數f必然要定義在這個要被返回內部函數中，而參數timeout只能定義在print_run_time()這個外部函數中了。此時print_run_time函數的定義應該是這樣的：

import time

def print_run_time(timeout=1):
    def decorator(f):
        def wrapper(*args, **kwargs):
            time_start = time.time()
            ret = f()
            time_end = time.time()
            func_name = f.__name__
            run_time = time_end - time_start
            print('%s run time is: %s' % (func_name, run_time))
            if run_time > timeout:
                print('PROBLEM'.rjust(30, '>'))
        return wrapper
    return decorator

此時的print_run_time已經不再是一個裝飾器函數，而是一個返回裝飾器函數的函數。func1使用這個函數語法糖應該是這樣的：

@print_run_time(timeout=2)
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

這個過程是這樣的：

decorator_func = print_run_time(timeout=2)

@decortator_func
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

函���調用：

func1('Decorator Test...')

輸出結果：

Decorator Test...
func1
func1 run time is: 1.0006623268127441

由於timeout是個默認參數（默認值為1），因此也可以不給它傳遞值，但是那對小括號不能省略：

@print_run_time()
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

此時的輸出結果中就會多打印一條錯誤信息：

Decorator Test...
func1
func1 run time is: 1.0006630420684814
>>>>>>>>>>>>>>>>>>>>>>>PROBLEM

4. 最終實現

看上去，上面已經把所有需求都實現了。不要擔心，還差一小步而已。此時我們可以打印一下被裝飾函數的__name__屬性看下：

print(func1.__name__)
print(func2.__name__)

輸出結果：

wrapper
wrapper

發現函數func1和func2的名字都變成了wrapper，當然其實還有一些其他屬性都變成了print_run_name的內部函數wrapper的屬性值。這顯然是不合適的，比如一些依賴函數簽名的代碼就可能會出現錯誤。解決這個問題也很簡單，就是給print_run_name的內部函數warpper應用一個裝飾器functools.wrapper就可以了(關於該裝飾器的說明請自行翻閱python文檔的funcstools模塊)：

def print_run_time(timeout=1):
    def decorator(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            time_start = time.time()
            ret = f(*args, **kwargs)
            time_end = time.time()
            func_name = f.__name__
            run_time = time_end - time_start
            print('%s run time is: %s' % (func_name, run_time))
            if run_time > timeout:
                print('PROBLEM'.rjust(30, '>'))
        return wrapper
    return decorator

此時再打印func1和func2的__name__屬性值看看：

print(func1.__name__)
print(func2.__name__)

輸出結果：

func1
func2

是的，這次沒問題了。

關於Python的裝飾器就先說到這裡，希望對大家有所幫助。