Python的装饰器

rs勿忘初心

已于 2022-11-26 16:14:22 修改

阅读量285

点赞数 1

分类专栏： # Python技术学习文章标签： python 装饰器装饰器作用

于 2021-11-03 23:20:23 首次发布

本文链接：https://blog.csdn.net/sinat_33718563/article/details/121098197

版权

Python技术学习专栏收录该内容

89 篇文章 10 订阅

订阅专栏

装饰器一直以来都是 Python 中很有用、很经典的一个 feature，在工程中的应用也十分广泛，比如日志、缓存等等的任务都会用到。

所谓的装饰器，其实就是通过装饰器函数，来修改原函数的一些功能，使得原函数不需要修改。

1.函数 -> 装饰器

必须掌握的函数的几个核心概念:

第一点，在 Python 中，函数是一等公民（first-class citizen），函数也是对象。我们可以把函数赋予变量，比如：

def func1(name):
    print("hello, {}".format(name))


get_name = func1
get_name('xiao ming')

上述示例中，把函数 func 赋予了变量 send_message，这样之后你调用 send_message，就相当于是调用函数 func()。

第二点，我们可以把函数当作参数，传入另一个函数中，比如下面这段代码：

def func1(name):
    return "hello, {}".format(name)


def func2(func，name):
    print(func(name))

    
func2(func1, 'xiao ming')

# 输出结果
hello, xiao ming

第三点，我们可以在函数里定义函数，也就是函数的嵌套。这里同样举了一个例子：

def func1(message):
    def func2(message):
        print('Got a message: {}'.format(message))
    return func2(message)

func1('xiao ming')

# 输出
Got a message: xiao ming

第四点，要知道，函数的返回值也可以是函数对象（闭包），比如下面这个例子：

def func1():
    def func2(message):
        print('Got a message: {}'.format(message))
    return func2

get_name = func1()
get_name('xiao ming')

# 输出
Got a message: xiao ming

2.装饰器

2.1 简单的装饰器

先来看一个装饰器的简单例子：


def my_decorator(func):
    def wrapper():
        print('wrapper of decorator')
        func()
    return wrapper

def greet():
    print('hello world')

greet = my_decorator(greet)
greet()

# 输出
wrapper of decorator
hello world

这段代码中，变量 greet 指向了内部函数 wrapper()，而内部函数 wrapper() 中又会调用原函数 greet()，因此，最后调用 greet() 时，就会先打印'wrapper of decorator'，然后输出'hello world'。

这里的函数 my_decorator() 就是一个装饰器，它把真正需要执行的函数 greet() 包裹在其中，并且改变了它的行为，但是原函数 greet() 不变。

事实上，上述代码在 Python 中有更简单、更优雅的表示：

def my_decorator(func):
    def wrapper():
        print('wrapper of decorator')
        func()
    return wrapper

# greet作为参数（函数形式）给到my_decorator
@my_decorator
def greet():
    print('hello world')

greet()

这里的@称之为语法糖，@my_decorator就相当于前面的greet=my_decorator(greet)语句，只不过更加简洁。因此，如果你的程序中有其它函数需要做类似的装饰，你只需在它们的上方加上@decorator就可以了，这样就大大提高了函数的重复利用和程序的可读性。

2.2 带有参数的装饰器

你或许会想到，如果原函数 greet() 中，有参数需要传递给装饰器怎么办？

def my_decorator(func):
    def wrapper(message):
        print('wrapper of decorator')
        func(message)
    return wrapper


@my_decorator
def greet(message):
    print(message)


greet('hello world')

# 输出
wrapper of decorator
hello world

不过，新的问题来了。如果我另外还有一个函数，也需要使用 my_decorator() 装饰器，但是这个新的函数有两个参数，又该怎么办呢？比如：

@my_decorator
def celebrate(name, message):
    ...

事实上，通常情况下，我们会把*args和**kwargs，作为装饰器内部函数 wrapper() 的参数。*args和**kwargs，表示接受任意数量和类型的参数，因此装饰器就可以写成下面的形式：

def my_decorator(func):
    def wrapper(*args, **kwargs):
        print('wrapper of decorator')
        func(*args, **kwargs)
    return wrapper

2.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


# 循环执行4次
@repeat(4)
def greet(message):
    print(message)


greet('hello world')

输出结果：

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

原函数还是原函数？还是之前的例子，我们试着打印出 greet() 函数的一些元信息：

greet.__name__
# 输出
'wrapper'

help(greet)
# 输出
Help on function wrapper in module __main__:

wrapper(*args, **kwargs)

你会发现，greet() 函数被装饰以后，它的元信息变了。元信息告诉我们“它不再是以前的那个 greet() 函数，而是被 wrapper() 函数取代了”。

为了解决这个问题，我们通常使用内置的装饰器@functools.wrap，它会帮助保留原函数的元信息（也就是将原函数的元信息，拷贝到对应的装饰器函数里）。

实际作用举例：多个函数被同一装饰器装饰，出错了，你就知道那个函数问题了。比如有两个函数A和B，分别加了装饰器C，并且装饰器C已加@functools.wraps(func)，那么函数A报错的时候，就能迅速定位哪个函数出错了。

import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('wrapper of decorator')
        func(*args, **kwargs)
    return wrapper
    
@my_decorator
def greet(message):
    print(message)

greet.__name__

# 输出
'greet'

2.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。

2.5 装饰器的嵌套

回顾刚刚讲的例子，基本都是一个装饰器的情况，但实际上，Python 也支持多个装饰器，比如写成下面这样的形式：

@decorator1
@decorator2
@decorator3
def func():
    ...

它的执行顺序从里到外，所以上面的语句也等效于下面这行代码：

decorator1(decorator2(decorator3(func)))

示例：

import functools


def my_decorator1(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('execute decorator1')
        func(*args, **kwargs)
    return wrapper


def my_decorator2(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('execute decorator2')
        func(*args, **kwargs)
    return wrapper


def my_decorator3(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('execute decorator3')
        func(*args, **kwargs)
    return wrapper


@my_decorator1
@my_decorator2
@my_decorator3
def greet(message):
    print(message)


greet('hello world')

运行结果：

execute decorator1
execute decorator2
execute decorator3
hello world

2.6 装饰器用法实例

（1）身份认证

首先是最常见的身份认证的应用。这个很容易理解，举个最常见的例子，你登录微信，需要输入用户名密码，然后点击确认，这样，服务器端便会查询你的用户名是否存在、是否和密码匹配等等。如果认证通过，你就可以顺利登录；如果不通过，就抛出异常并提示你登录失败。

再比如一些网站，你不登录也可以浏览内容，但如果你想要发布文章或留言，在点击发布时，服务器端便会查询你是否登录。如果没有登录，就不允许这项操作等等。

大概的代码示例：

import functools

def authenticate(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        request = args[0]
        if check_user_logged_in(request): # 如果用户处于登录状态
            return func(*args, **kwargs)  # 执行函数post_comment() 
        else:
            raise Exception('Authentication failed')
    return wrapper
    
@authenticate
def post_comment(request, ...)
    ...

这段代码中定义了装饰器 authenticate；而函数 post_comment()，则表示发表用户对某篇文章的评论。每次调用这个函数前，都会先检查用户是否处于登录状态，如果是登录状态，则允许这项操作；如果没有登录，则不允许。

（2）日志记录

日志记录同样是很常见的一个案例。在实际工作中，如果你怀疑某些函数的耗时过长，导致整个系统的 latency（延迟）增加，所以想在线上测试某些函数的执行时间，那么，装饰器就是一种很常用的手段。

我们通常用下面的方法来表示：

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_time():
    time.sleep(3)


calculate_time()

输出：

calculate_time took 3004.4800929899793 ms

这里，装饰器 log_execution_time 记录某个函数的运行时间，并返回其执行结果。如果你想计算任何函数的执行时间，在这个函数上方加上@log_execution_time即可。

（3）合理性检查

再来看今天要讲的第三个应用，输入合理性检查。在大型公司的机器学习框架中，我们调用机器集群进行模型训练前，往往会用装饰器对其输入（往往是很长的 JSON 文件）进行合理性检查。这样就可以大大避免，输入不正确对机器造成的巨大开销。

它的写法往往是下面的格式：

import functools

def validation_check(input):
    @functools.wraps(func)
    def wrapper(*args, **kwargs): 
        ... # 检查输入是否合法
    
@validation_check
def neural_network_training(param1, param2, ...):
    ...

其实在工作中，很多情况下都会出现输入不合理的现象。因为我们调用的训练模型往往很复杂，输入的文件有成千上万行，很多时候确实也很难发现。

试想一下，如果没有输入的合理性检查，很容易出现“模型训练了好几个小时后，系统却报错说输入的一个参数不对，成果付之一炬”的现象。这样的“惨案”，大大减缓了开发效率，也对机器资源造成了巨大浪费。

（4）缓存

最后，我们来看缓存方面的应用。关于缓存装饰器的用法，其实十分常见，这里我以 Python 内置的 LRU cache 为例来说明（如果你不了解 LRU cache，可以点击链接自行查阅）。

LRU cache，在 Python 中的表示形式是@lru_cache。@lru_cache会缓存进程中的函数参数和结果，当缓存满了以后，会删除 least recenly used 的数据。

正确使用缓存装饰器，往往能极大地提高程序运行效率。为什么呢？我举一个常见的例子来说明。

大型公司服务器端的代码中往往存在很多关于设备的检查，比如你使用的设备是安卓还是iPhone，版本号是多少。这其中的一个原因，就是一些新的 feature，往往只在某些特定的手机系统或版本上才有（比如 Android v200+）。

这样一来，我们通常使用缓存装饰器，来包裹这些检查函数，避免其被反复调用，进而提高程序运行效率，比如写成下面这样：


@lru_cache
def check(param1, param2, ...) # 检查用户设备类型，版本号等等
    ...

常见问题：

1、AssertionError: View function mapping is overwriting an existing endpoint function:.wrapper

存在的问题用法：

def host_whiltlist(func):
    """
    # 机器白名单检查装饰器
    """
    def wrapper():
        """
        # 
        """
        # 请求来源ip
        print("hello, world!")
    return wrapper


@one_app.route('/api/hello1', methods=['POST'])
@host_whiltlist
def hello1():
    print("hello, wolrd!")


@one_app.route('/api/hello2', methods=['POST'])
@host_whiltlist
def hello2():
    print("hello, wolrd!")


....

装饰器函数需要添加：@wraps(func)，参考：AssertionError: View function mapping is overwriting an existing endpoint function: login

def host_whiltlist(func):
    """
    # 机器白名单检查装饰器
    """
    @wraps(func)
    def wrapper():
        """
        # 
        """
        # 请求来源ip
        print("hello, world!")
    return wrapper