Python的装饰器

在Python中使用装饰器可以在不修改代码的前提下为已有的函数添加新功能. 例如计算函数的运行时间, 打印函数的运行日志, 缓存数据, 用户授权等.

为什么需要装饰器

为什么需要把某个功能用装饰器的方式实现? 换句话说, 为什么不直接在函数中实现某个功能? 或者为什么不单独用一个类(class)来实现某个功能, 然后在函数中调用?

以提供缓存功能的装饰器为例. 假如我们的业务代码负责向数据库读取数据, 然后返回给用户. 为了缩短响应时间, 我们可以把数据放在缓存, 从而降低IO. 站在业务的角度, 缓存功能实际上与我们需要实现的业务功能无关. 为了保持业务代码的简洁, 我们可以把通用的与业务逻辑无关的功能用装饰器来实现, 即面向切面编程的思想.

普通装饰器

我们想实现一个计时的装饰器timer. 使用效果如下所示:

# main.py
import time

@timer
def test():
	time.sleep(1.0)

if __name__ == '__main__':
	test()

time cost: 1.0013980865478516

装饰器timer的输入参数是一个函数对象, 返回的结果也是一个函数对象. 返回的函数是一个“附加了新功能的”函数. 在上面的例子中, 装饰器@timer实际上是一个执行如下操作的语法糖.

wrapper_func = timer(test)
wrapper_func()

完整代码

import time


# 装饰器函数
# 输入是一个函数, 返回一个装饰过的函数wrapper
def timer(func):
	
	# 调用func并计时
    def wrapper(*args, **kwargs):
        t1 = time.time()
        func(*args, **kwargs)
        print("time cost:", time.time() - t1)

    return wrapper

@timer
def test():
	time.sleep(1.0)

if __name__ == '__main__':
	test()

装饰后的函数名

在上述例子中, 如果我们打印使用装饰器之后的函数test的名称(结果是wrapper).

>>> print(test.__name__)
wrapper

在一些情况下, 我们希望保持原来函数的名称. 这个时候我们可以利用Python自带的装饰器@wraps(func)来装饰wrapper, 从而保持函数名称不变.

import time
from functools import wraps


# 装饰器函数
# 输入是一个函数, 返回一个装饰过的函数wrapper
def timer(func):
	
	# 调用func并计时
	@wraps(func)  # 保持func的函数名
    def wrapper(*args, **kwargs):
        t1 = time.time()
        func(*args, **kwargs)
        print("time cost:", time.time() - t1)

    return wrapper

带参数的装饰器

下面我们要把计时器的功能稍微扩充一下. 实现一个带参数的计时器@timer_unit(unit), 其中参数unit代表了计时的单位: 's' -- 秒; 'ms' -- 毫秒. 前文提到装饰器函数的输入必须是一个函数对象, 否则如何接收函数对象? 注意到函数名加括号timer_unit(unit='s')代表执行一个函数, 因此我们只需要让函数timer_unit(unit='s')的执行结果返回一个普通的装饰器即可.

# 带参数的装饰器
def timer_unit(unit='s'):
	
	# 普通装饰器
    def decorator(func):
		
		# 调用func并计时
		@wraps(func)  # 保持func的函数名
        def wrapper(*args, **kwargs):
            t1 = time.time()
            func(*args, **kwargs)
            multiplier = 1
            if unit == 'ms':
                multiplier = 1000
            print("time cost (unit: %s):" % unit, (time.time() - t1) * multiplier)

        return wrapper
        
	# 返回普通装饰器
    return decorator

装饰器类

假如我们要给上述装饰器增加新的计时单位'mus: 微秒'. 简单的方法是在wrapper函数中增加相应的逻辑. 这样做的缺点是随着我们需要的功越来越复杂, 所有的逻辑必须在一个函数里实现, 这样不利于协作和维护. 因此, 在功能复杂的情况下我们考虑用类的方式来实现装饰器的功能.

import time
from functools import wraps


class Timer(object):
    """ 一个用来计时的装饰器类
    需要实现__call__方法, 可以像函数一样调用实例
    """

    def __init__(self):
        self._multiplier = 1  # 秒
        self._unit = None  # 's'-秒; 'ms'-毫秒; 'mus' - 微秒

    def _get_multiplier(self):
        if self._unit == 'ms':  # 毫秒
            self._multiplier = 1000
        if self._unit == 'mus':  # 微秒
            self._multiplier = 1000000
        return self._multiplier

    def __call__(self, unit='s'):

        self._unit = unit

        def decorator(func):
        
        	@wraps(func)
            def wrapper(*args, **kwargs):
                t1 = time.time()
                func(*args, **kwargs)
                print("time cost (unit: %s):" % unit, (time.time() - t1) * self._get_multiplier())

            return wrapper

        return decorator


# 装饰器函数
timer_class = Timer()


@timer_class(unit='mus')
def test():
    time.sleep(1)


if __name__ == '__main__':
    test()