Python装饰器

Posted 2022-06-08 詩和遠方

装饰器（decorator）是python特有的语法特性，通过函数封装扩展原函数的功能。
下面一步步通过例子说明它的作用和用法。

函数中定义函数

def parent():
    print("Printing from the parent() function")

    def first_child():
        print("Printing from the first_child() function")

    def second_child():
        print("Printing from the second_child() function")

    second_child()
    first_child()

parent()

Printing from the parent() function
Printing from the second_child() function
Printing from the first_child() function

函数引用（类似函数指针）

函数也是一个对象，可以赋值给变量，传引用，相当于起了一个别名

my_print = print
my_print('Hello World')

Hello World

返回函数引用的函数

既然函数也是一个对象，那就可以被函数当做返回值，如下：

def parent(num):
    def first_child():
        return "Hi,I am Melody"
    def second_child():
        return "Call me Susan"
    
    if num%2 == 1:
        return first_child
    else:
        return second_child

first = parent(3)
second = parent(4)
print(first())
print(second())

Hi,I am Melody
Call me Susan

以上代码根据参数不同，执行不同的函数

简单装饰器

def my_decorator(func):
    def wrapper():
        print(f"do something before func.__name__ is called")
        func()
        print(f"do something after func.__name__ is called")
    return wrapper

def say_hello():
    print("Hello")

say_hello = my_decorator(say_hello)
say_hello()

do something before say_hello is called
Hello
do something after say_hello is called

以上例子说明：装饰器就是一种特殊的函数，接收函数引用作为参数，对函数进行封装，然后改变或改进它的行为，最后返回一个函数引用。

装饰器语法糖

def my_decorator(func):
    def wrapper():
        print(f"do something before func.__name__ is called")
        func()
        print(f"do something after func.__name__ is called")
    return wrapper

@my_decorator
def say_hello():
    print("Hello")

say_hello()

do something before say_hello is called
Hello
do something after say_hello is called

上面代码中的@my_decorator相当于【在函数定义之后加了一句say_hello = my_decorator(say_hello)】的简写。
这里需要注意，my_decorator不一定返回另一个函数的引用，它也可以执行一些额外的代码后，还是返回被装饰函数本身。

装饰器复用

可以将一些常用的函数包装器放在一个文件中方便后续进行复用

# 文件 decorators.py
def do_twice(func):
    def wrapper_do_twice():
        func()
        func()
    return wrapper_do_twice

# 引用上述文件中的方法
from decorators import do_twice

@do_twice
def say_hello():
    print("Hello")
    
say_hello()

Hello
Hello

带参数的装饰器

from decorators import do_twice

@do_twice
def say_hello_to(name):
    print(f"Hello, name")
    
say_hello_to("Susan")

执行后报错：TypeError: wrapper_do_twice() takes 0 positional arguments but 1 was given

那要如何避免上述问题呢，答案是用*args和**wargs
在decorators.py中加入如下新函数：

# 文件 decorators.py
def do_twice_parms(func):
    def wrapper_do_twice(*args,**kwargs):
        func(*args,**kwargs)
        func(*args,**kwargs)
    return wrapper_do_twice

from decorators import do_twice_parms

@do_twice_parms
def say_hello_to(name):
    print(f"Hello, name")
    
say_hello_to("Susan")

Hello, Susan
Hello, Susan

被装饰函数的元数据更新

help(say_hello_to)

Help on function say_hello_to in module __main__:
say_hello_to(name)

可以发现查看say_hello_to函数相关信息（元数据）时，指示为装饰器函数内的wrapper_do_twice，显然这不是我们想要的，如何改进呢？

本质上需要更新函数的__name__、__doc__、__module__属性，让其指向原函数，python有个惯用手法来解决此问题，那就是用functools模块的wrap方法：

import functools

def do_twice_parms(func):
    @functools.wraps(func)
    def wrapper_do_twice(*args,**kwargs):
        func(*args,**kwargs)
        func(*args,**kwargs)
    return wrapper_do_twice

@do_twice_parms
def say_hello_to(name):
    print(f"Hello, name")
    
say_hello_to("Susan")

Hello, Susan
Hello, Susan

再次查看say_hello_to的帮助信息，已指向原函数：

help(say_hello_to)

Help on function say_hello_to in module __main__:
say_hello_to(name)

装饰器常用模板

从以上例子我们可以总结出装饰器函数的常用模板：

import functools

def decorator(func):
    @functools.wraps(func)
    def wrapper_decorator(*args, **kwargs):
        # Do something before
        value = func(*args, **kwargs)
        # Do something after
        return value
    return wrapper_decorator

实际应用：函数执行时间

下面举个实际应用的例子，调用函数时，同时打印执行的时间

import functools
import time

def timer(func):
    """打印函数执行的时长"""
    @functools.wraps(func)
    def wrapper_timer(*args, **kwargs):
        start_time = time.time()
        value = func(*args, **kwargs)
        end_time = time.time()
        print(f"Finished func.__name__ in end_time-start_time:.3f seconds")
        return value
    return wrapper_timer

@timer
def do_some_loops(num_times):
    for _ in range(num_times):
        sum([i*i for i in range(10000)])

do_some_loops(9)
do_some_loops(99)
do_some_loops(999)

Finished do_some_loops in 0.020 seconds
Finished do_some_loops in 0.165 seconds
Finished do_some_loops in 1.234 seconds

参考资源：
https://realpython.com/primer-on-python-decorators/
https://docs.python.org/3/library/functools.html#functools.wraps

开发者涨薪指南 48位大咖的思考法则、工作方式、逻辑体系