【Python測試開發】裝飾器的應用

⽽實際⼯作中，裝飾器通常運⽤在身份認證、⽇志記錄、輸⼊合理性檢查等多個領域中。合理使⽤裝飾器，往往能極⼤地提⾼程序的可讀性以及運⾏效率。

所謂的裝飾器，其實就是通過裝飾器函數或者裝飾器類，來修改原函數的功能，使得原函數不需要修改就具備新特性的機制。

1.一個簡單的裝飾器

我們可以先來看一個裝飾器的簡單例子：

import functools


def makebold(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return "<b>" + fn(*args, **kwargs) + "</b>"

    return wrapper

@makebold   
def greet(message): 
    return message
    
@makebold   
def greet2(): 
    return "hello"
    
print(greet.__name__)
print(greet('hello world'))
print(greet2())

上面代碼的輸出：

greet
<b>hello world</b>
<b>hello</b>

這段代碼中，makehold就是一個裝飾器函數，內部函數wrapper()中調⽤了原函數fn(*args, **kwargs)，在原函數返回值的前後分別拼上了<b>和</b>，從而改變了原函數的行爲。裝飾器函數返回內部函數wrapper。

用@makebold語法去修飾greet函數，這樣原函數 greet() 不需要任何變化，在調用greet()時，就會在原來返回值message前後增加<b>和</b>字符串了。

這⾥的@被稱之爲語法糖，@makebold相當於前⾯的greet=makebold(greet)語句。因此，如果你的程序中有其它函數需要做類似的裝飾，你只需在它們的上⽅加上@makebold就可以了，這樣就⼤⼤提⾼了函數的重複利⽤和程序的可讀性。

內部函數wrapper使⽤內置的裝飾器@functools.wrap裝飾，它會幫助保留原函數的元信息（也就是將原函數的元信息，拷⻉到對應的裝飾器函數⾥）。如果不用裝飾器@functools.wrap，那麼greet.__name__的值將是wrapper。

2.裝飾器的嵌套

Python支持對一個函數同時應用多個裝飾器，多個裝飾器，主要是要關注裝飾器執行順序問題：

import functools


def makebold(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return "<b>" + fn(*args, **kwargs) + "</b>"

    return wrapper


def makeitalic(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return "<i>" + fn(*args, **kwargs) + "</i>"

    return wrapper

@makebold
@makeitalic
def hello(s):
    return s

@makeitalic
@makebold
def world(s):
    return s

print(hello('hello'))
print(world('world'))

上面代碼的輸出：

<b><i>hello</i></b>
<i><b>world</b></i>

上面代碼中，hello()和world()均被兩個裝飾器裝飾了，只是裝飾器的順序不同。從代碼的輸出中看到，因爲裝飾器的順序不同，輸出也是不同的。

裝飾器的執行順序是從下往上。hello()函數裝飾後等效於makebold(makeitalic(hello))，world()函數裝飾後等效於函數裝飾後等效於makeitalic(makebold(hello))。

3. 裝飾器本身帶參數

裝飾器可以接受原函數任意類型和數量的參數，除此之外，它還可以接受⾃⼰定義的參數。舉個例⼦，⽐如我想要定義⼀個參數，來表示裝飾器內部函數被執⾏的次數，那麼就可以寫成下⾯
這種形式：

def repeat(num):
    def my_decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(num):
                print('wrapper of decorator')
                func(*args, **kwargs)
        return wrapper
    return my_decorator


@repeat(3)
def greet(message):
    print(message)

greet('hello world')

上面代碼輸出：

wrapper of decorator
hello world
wrapper of decorator
hello world
wrapper of decorator
hello world

從輸出上看，greet函數被執行了3次。仔細觀察上面裝飾器，是在原有不帶參數裝飾器函數外邊又套了一層。

4. 類裝飾器

類也可以作爲裝飾器。類裝飾器主要依賴於函數__call_()，每當你調⽤⼀個類的示例時，函數__call__()就會被執⾏⼀次。看個簡單的例子：

class Count:
    def __init__(self, func):
        self.func = func
        self.num_calls = 0

    def __call__(self, *args, **kwargs):
        self.num_calls += 1
        print('num of calls is: {}'.format(self.num_calls))
        return self.func(*args, **kwargs)

@Count
def example():
    print("hello world")

example()
example()

上面代碼的輸出：

num of calls is: 1
hello world
num of calls is: 2
hello world

Count類是一個裝飾器類，初始化時傳⼊原函數func()，⽽__call__()函數表示讓變量num_calls ⾃增 1，然後打印，並且調⽤原函數。因此，在我們第⼀次調⽤函數 example() 時，num_calls 的值是 1，⽽在第⼆次調⽤時，它的值變成了2。

再舉一個例子：

class Foo(object):
    def __init__(self):
        pass

    def __call__(self, func):
        def _call(*args, **kw):
            print('class decorator runing')
            return func(*args, **kw)

        return _call


class Bar(object):
    @Foo()
    def bar(self, test, ids):
        print(test, ids)


Bar().bar('aa', 'ids')

我們經常會用到的比如說給多線程函數加鎖來保證共享數據操作的完整性。線程鎖裝飾器如下：

import threading
class StandardOut(object): 
   """ 
   線程鎖裝飾器
   """ 
   def __init__(self):
       self.thread_lock = threading.Lock()
    
   def __call__(self,func):
       def _call(*args,**kw): 
           self.thread_lock.acquire()
           func(*args,**kw)
           self.thread_lock.release()
       return _call

5.裝飾器的實際應用

5.1 身份認證

⾸先是最常⻅的身份認證的應⽤。⽐如⼀些⽹站，你不登錄也可以瀏覽內容，但如果你想要發佈⽂章或留⾔，在點擊發布時，服務器端便會查詢你是否登錄。如果沒有登錄，就不允許這項操作。

我們來看一個代碼示例：

import functools

def check_user_logged_in(request):
    """
    判斷用戶是否處於登錄狀態。
    """
    pass

def authenticate(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if check_user_logged_in(args[0]): # 如果用戶處於登錄狀態
            return func(*args, **kwargs) # 執行函數 post_comment() 
        else:
            raise Exception('Authentication failed')
    return wrapper
    
@authenticate
def post_comment(request):
    pass

authenticate是一個裝飾器，在內部函數wrapper中，判斷用戶是否處於登錄狀態，如果是則執行被裝飾函數func，否則拋出異常"Authentication failed"。

post_comment()是一個發表評論的函數，使用authenticate對其進行裝飾，每次調⽤這個函數前，都會先檢查⽤戶是否處於登錄狀態，如果是登錄狀態，則允許這項操作；如果沒有登錄，則不允許。

5.2 ⽇志記錄

⽇志記錄同樣是很常⻅的⼀個案例。在實際⼯作中，想在測試某些函數的執⾏時間，那麼，裝飾器就是⼀種很常⽤的⼿段。

我們通常⽤下⾯的⽅法來表示：

import time
import functools


def log_execution_time(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        res = func(*args, **kwargs)
        end = time.perf_counter()
        print('{} took {} ms'.format(func.__name__, (end - start) * 1000))
        return res

    return wrapper


@log_execution_time
def calculate_similarity():
    time.sleep(3)


calculate_similarity()

上面代碼輸出：

calculate_similarity took 3001.337694 ms

這⾥，裝飾器 log_execution_time 記錄某個函數的運⾏時間，並返回其執⾏結果。如果你想計算任何函數的執⾏時間，在這個函數上⽅加上@log_execution_time即可。

5.3 參數的合理性檢查

一些通用的參數檢查可以放到一個裝飾器函數中，然後使用裝飾器去裝飾需要做參數檢查的函數：

import functools


def validation_check(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if all([isinstance(i, str) for i in args]):
            return func(*args, **kwargs)
        else:
            raise Exception('位置參數不全是擦字符串類型')

    return wrapper


@validation_check
def neural_network_training(param1, param2):
    print(param1, param2)


neural_network_training("123", "a") # 輸出123 a
neural_network_training("123", 123) # 輸出異常

這個例子是對neural_network_training函數的參數做檢查，當位置參數全是字符串類型是才執行，否則不執行。

總結

當我們希望原始函數保持不變的情況下，對其增加新的特性時，就可以使用裝飾器。裝飾器的應用範圍很廣，比如身份認證、⽇志記錄、輸⼊合理性檢查、路由等多個場景。

裝飾器可以有兩種類型，一種是函數裝飾器，一種是類裝飾器，寫成類的話，優點是程序的分解度更加高，具體用類裝飾器和函數裝飾器，視情況而定，二者本質是一樣的。