[Python] Decorator 裝飾器 (2) - wraps

上篇文章：
[Python] Decorator 裝飾器

在使用 decorator時，會使得原本被包裹的 function 的 attribute資料遺失。

import time

def measuretime(func):

    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print (func.__name__, end - start)
        return result

    return wrapper

@measuretime
def count_down(n):
    '''Counts down from n'''
    while n > 0:
        n -= 1

if __name__ == '__main__':
    print (count_down.__name__)
    print (count_down.__doc__)

Output：

wrapper
None

可使用 functools的 decorator wraps來避免這種情況：

import time
from functools import wraps

def measuretime(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

@measuretime
def count_down(n):
    '''Counts down from n'''
    while n > 0:
        n -= 1

if __name__ == '__main__':
    print (count_down.__name__)
    print (count_down.__doc__)

Output：

count_down
Counts down from n

Unwrapping

wraps提供了另一個有用的功能，可以使已經套用過 decorator的 function回復成套用前的
function。透過 __wrapped__這個屬性來使用原本的 function。

import time
from functools import wraps

def measuretime(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

@measuretime
def count_down(n):
    '''Counts down from n'''
    print (f'Counts down from {n}')
    while n > 0:
        n -= 1

if __name__ == '__main__':
    count_down(500000)
    org_count_down = count_down.__wrapped__
    org_count_down(500000)

Output：

Counts down from 500000
count_down 0.021966218948364258
Counts down from 500000