Python學習筆記 - 裝飾器

MADAO是不會開花的發表於2019-01-08
Python筆記

有時候我們希望給一個函式新增一些功能，但是又不想對原有函式進行修改，那麼這時候就用到了裝飾器或者說是裝飾模式。

例子：

# 1. 有一個求和函式sum
def sum(a, b):
    return a + b
    
# 2. 如果想要知道這個函式的執行時間，如何實現?

# 實現一
import time

def sum(a, b):
    return a + b
start = time.time()
sum(5, 6)
end = time.time()

print('sum函式的執行時間大約是：%f秒' % (end - start))

# 3. 如果還有一個求差函式sub，要知道這個函式的執行時間，如何實現?
'''
按照實現獲取sum函式執行時間的思路，完全可以照搬上面的程式碼，之前說過有重複
地方儘量封裝成函式。
'''

# 實現二

import time


def print_execution_time(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print('%s函式的執行時間大約是：%f秒' % (func.__name__, end - start))
        return result
    return wrapper


def sum(a, b):
    return a + b


def sub(a, b):
    return a - b


print_execution_time(sum)(10, 20)
print_execution_time(sub)(10, 20)

複製程式碼

這樣就實現了一個函式用於獲取某個函式的執行時間， 這種模式就是裝飾模式，有點像閉包，只不過傳遞的引數 換成了傳遞函式。

但是這樣做有點問題，就是原本函式的呼叫方式改變了。而且函式主要的目的也變得 很模糊，原函式主要是用於計算，改寫成這樣後就好像獲取執行時間才是主要目的。

Python提供了一個語法糖 - 裝飾器@。通過這個可以寫出更為簡潔的程式碼。例子：

# 上面例子使用Python提供的裝飾器改寫：

import time


def print_execution_time(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print('%s函式的執行時間大約是：%f秒' % (func.__name__, end - start))
        return result
    return wrapper

@print_execution_time  # 新增
def sum(a, b):
    return a + b

@print_execution_time  # 新增
def sub(a, b):
    return a - b

sum(10, 10)
sub(1, 3)
複製程式碼

結果：

Python學習筆記 - 裝飾器

用了裝飾器之後，呼叫和以前完全一樣。只是在定義函式的時候加上@裝飾函式名這種語法，原本函式就獲得了裝飾函式中新增的功能。

上面例子中我們都用了func.__name__，這句程式碼的意思是獲取函式的名字。但是用了上面裝飾器之後我們再看一下sum和sub函式的__name__屬性：

print('sum的__name__是：%s' % (sum.__name__))
print('sub的__name__是：%s' % (sub.__name__))
複製程式碼

Python學習筆記 - 裝飾器

變了，全部變成wrpper了。這是因為使用@print_execution_time的時候實際上等於 sum = print_execution_time(sum)，得到結果就是wrapper函式。我們可以使用funcools模組中的wrap方法來改正這個問題：

import time
from functools import wraps # 新增


def print_execution_time(func):
    @wraps(func)                  # 新增
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print('%s函式的執行時間大約是：%f秒' % (func.__name__, end - start))
        return result
    return wrapper


@print_execution_time
def sum(a, b):
    return a + b


@print_execution_time
def sub(a, b):
    return a - b


sum(10, 10)
sub(1, 3)

print('sum的__name__是：%s' % (sum.__name__))  # sum的__name__是：sum
print('sub的__name__是：%s' % (sub.__name__))  # sub的__name__是：sub
複製程式碼

裝飾器本身也是可以帶引數的，帶引數的裝飾器寫法如下：

import time
from functools import wraps


def print_execution_time(text):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            end = time.time()
            print('%s函式的執行時間大約是：%f秒' % (func.__name__, end - start))
            print(text)
            return result
        return wrapper
    return decorator 


@print_execution_time('今天天氣不錯')
def sum(a, b):
    return a + b


@print_execution_time('今天天氣不好')
def sub(a, b):
    return a - b


sum(10, 10)
sub(1, 3)

print('sum的__name__是：%s' % (sum.__name__))
print('sub的__name__是：%s' % (sub.__name__))

複製程式碼

結果：

Python學習筆記 - 裝飾器

寫法就是在原有的裝飾器上再包裹一層函式。

@print_execution_time('今天天氣不錯') 相當於

sum = print_execution_time('今天天氣不錯')(sum)

總結下兩種寫法：

  1. 不帶引數的裝飾器
from functools import wraps


def decorator(func):
    @wraps(func)
    def wrapper(*args, **wkargs):
        # 寫新增功能
        result = func(*args, **wkargs)
        return result
    return wrapper   
複製程式碼
  1. 帶引數的裝飾器
from functools import wraps


from functools import wraps
def decorator_wrapper(*args, **wkargs):
    def decorator(func):
        @wraps(func)
        def wrapper(*sub_args, **sub_wkargs):
            # 寫新增功能
            result = func(*sub_args, **sub_wkargs)
            return result
        return wrapper
    return decorator
複製程式碼

相關文章