python装饰器

418 阅读5分钟

什么是装饰器

装饰器本质上是一个Python函数,它可以让其他函数在不需要做任何代码变动的前提下增加额外功能,装饰器的返回值也是一个函数对象。它经常用于有切面需求的场景,比如:插入日志、性能测试、事务处理、缓存、权限校验等场景。装饰器是解决这类问题的绝佳设计,有了装饰器,我们就可以抽离出大量与函数功能本身无关的雷同代码并继续重用。

概括的讲,装饰器的作用就是为已经存在的函数或对象添加额外的功能

如何定义装饰器

Python的装饰器是一个固定的函数接口形式。它要求你的装饰器函数(或装饰器类)必须要接受一个函数并返回一个函数:

def wrapper(func1): # 接受一个callable对象
    return func2    # 返回一个对象,一般为函数

@wrapper #使用@语法调用
def target_func(args):
    pass

如果被装饰的函数需要传入参数的话,可以指定装饰器函数wrapper接受和原函数一样的参数

def debug(func):
    def wrapper(something): # 指定一模一样的参数
        print ("[DEBUG]: enter{} ()".format(func.__name__))
	return func(something)
    return wrapper

@debug
def say_hello(something):
    print("hello {}!".format(something))

say_hello("ly")

不过函数有千千万,并不能保证所有函数的参数都一致。所以可以使用Python提供的可变参数*args和关键字参数**kwargs,有了这两个参数,装饰器就可以用于任意目标函数了。

def debug(func):
    def wrapper(*args, **kwargs):
        print ("[DEBUG]: enter{} ()".format(func.__name__))
	return func(*args, **kwargs)
    return wrapper

@debug
def say_hello(something):
    print("hello {}!".format(something))

say_hello("ly")

带参数的装饰器

假设我们前文的装饰器需要完成的功能不仅仅时能在进入某个函数后打出log信息,而且还需指定log的级别,那么装饰器就会是这样的。

def logging(level):
    def wrapper(func):
	def inner_wrapper(*args, **kwargs):
            print("[{level}]: enter function {func}()".format(level=level, func=func.__name__))
            return func(*args, **kwargs)
	return inner_wrapper
    return wrapper

@logging(level='INFO')
# 如果没有使用@语法,等同于 say=logging(level='INFO')(say)
def say(something):
    print("say {}!".format(something))

@logging(level='DEBUG')
def do(something):
    print("do {}...".format(something))

if __name__ == '__main__':
    say('hello')
    do("my work")

# [INFO]: enter function say()
# say hello!
# [DEBUG]: enter function do()
# do my work...

原函数还是原函数吗?

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

print(greet.__name__)

#输出:warpper

可以看到,greet()函数被装饰以后,它的元信息变了。元信息告诉我们“它不再是以前的那个greet()函数,而是被wrapper()函数取代了”

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

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'

类装饰器

类也可以作为装饰器。类装饰器主要依赖于函数__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()
example()


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

这里,我们定义了类Count,初始化时传入原函数func(),而__call__()函数表示让变量num_calls自增1,然后打印,并且调用原函数。因此,在我们第一次调用函数example()时,num_calls的值是1,而在第二次调用时,它的值变成了2。

装饰器的嵌套

python也支持多个装饰器,例

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

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

decorator1(decorator2(decorator3(func)))

装饰器用法实例

身份认证

比如一些网站,不登录也可以浏览内容,但如果你想要发布文章或留言,在点击发布时,就会查询是否登录

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)
	else:
            raise Exception('Authentication failed')
    return wrapper

@authenticate
def post_comment(request, ...):
	...

日志记录

在实际工作中,如果你怀疑某些函数的耗时过长,导致整个系统的 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_similarity(items):
    ...

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

输入合理性检查

使用装饰器对输入进行合理性检查,这样就可以大大避免,输入不正确对机器造成的巨大开销。

import functools

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

缓存

python内置的LRU cache,表示形式为@lru_cache。@lru_cache会缓存进程中的函数参数和结果,当缓存满了以后,会删除least recently used的数据。正确使用缓存装饰器,往往能极大的提高程序运行效率。

例如:大型公司服务器端的代码中往往存在很多关于设备的检查,比如你使用的设备是安卓还是 iPhone,版本号是多少。这其中的一个原因,就是一些新的 feature,往往只在某些特定的手机系统或版本上才有(比如 Android v200+)。这样一来,我们通常使用缓存装饰器,来包裹这些检查函数,避免其被反复调用,进而提高程序运行效率,比如写成下面这样:

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