博客

1 为什么需要装饰器

Leader让小A写两个数字相加和相减的函数，小A很快就写完了：

def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    result = add(1, 2)
    print(result)

    result = sub(5, 4)
    print(result)

# 输出：
# 3
# 1

Leader让小A添加上统计函数的运行时长的功能, 小A直接在调用函数时加上了时长的计算:

import time


def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    start = time.time()
    result_1 = add(1, 2)
    end = time.time()
    print('result: %d' % result_1)
    print('time taken %f' % (end - start))

    start = time.time()
    result_2 = sub(5, 4)
    end = time.time()
    print('result: %d' % result_2)
    print('time taken %f' % (end - start))

# 输出:
# result: 3
# time taken 0.000000
# result: 1
# time taken 0.000000

Leader看了，说每次调用函数岂不是要写很多重复代码吗。小A进行了优化：

import time


def add(x, y):
    start = time.time()
    rv = x + y
    end = time.time()
    print('time taken %f' % (end - start))
    return rv


def sub(x, y):
    start = time.time()
    rv = x - y
    end = time.time()
    print('time taken %f' % (end - start))
    return rv


if __name__ == '__main__':
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    result_2 = sub(5, 4)
    print('result: %d' % result_2)
    
# 输出
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

这种方法肯定比前一种要好。但是当我们有多个函数时，那么这似乎就不方便了。

小A又定义了一个计时的函数并包装其他函数，然后返回包装后的函数：

import time


def time_taken(func):
    def inner(*args, **kwargs):
        start = time.time()
        rv = func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        return rv

    return inner


def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    add = time_taken(add)  # 将函数作为参数传给另一个函数
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    sub = time_taken(sub)
    result_2 = sub(5, 4)
    print('result: %d' % result_2)

# 输出：
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

Leader说上面的解决方案以及非常接近装饰器的思想了，小A查了一下装饰器的用法，加入装饰器后代码果然变得很优雅。

import time


def time_taken(func):
    def inner(*args, **kwargs):
        start = time.time()
        rv = func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        return rv
    return inner


@time_taken
# @time_taken等价于add = time_taken(add)
def add(x, y):
    return x + y


@time_taken
def sub(x, y):
    return x - y


if __name__ == '__main__':
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    result_2 = sub(5, 4)
    print('result: %d' % result_2)

# 输出：
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

2 装饰器

2.1 装饰器的定义

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

2.2 给某个函数加上多个装饰器

import time


def logging(func):
    print('logging start')

    def inner(*args, **kwargs):
        print('logging.inner start')
        rv = func(*args, **kwargs)
        print('logging.inner end')
        return rv
    return inner


def time_taken(func):
    print('time_taken start')

    def inner(*args, **kwargs):
        print('time_taken.inner start')
        start = time.time()
        rv = func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        print('time_taken.inner end')
        return rv
    return inner


@logging
@time_taken
# 等价于 logging(time_taken(add))
def add(x, y):
    return x + y


if __name__ == '__main__':
    result = add(1, 2)
    print('result: %d' % result)

# 输出：
# time_taken start
# logging start
# logging.inner start
# time_taken.inner start
# time taken 0.000000
# time_taken.inner end
# logging.inner end
# result: 3

2.3 带参数的装饰器

2.3.1 闭包

在讲带参数的装饰器前，先简要介绍一下

闭包

的概念。如果在一个函数的内部定义了另一个函数，外部的我们叫它为外函数，内部的我们叫它为内函数。

闭包

： 在一个外函数中定义了一个内函数，内函数里运用了外函数的临时变量，并且外函数的返回值是内函数的引用。这样就构成了一个闭包。

def f1():
    n = 999

    def f2():
        print(n)

    return f2


if __name__ == '__main__':
    result = f1()
    result()

# 输出：
# 999

函数 f2 被包括在函数 f1 内部，这时 f1 内部的所有局部变量，对 f2 都是可见的。 既然 f2 可以读取 f1 中的局部变量，那么只要把 f2 作为返回值，我们不就可以在 f1 外部读取它的内部变量了吗? 上一部分代码中的

f2函数

，就是闭包 。

2.3.2 带参数的装饰器


def logging(func)

是一个装饰器，如果装饰器需要参数，需要通过闭包来实现，即在其外面再定义一个外函数

def mylog(type)

，将参数

type

作为外函数的的变量传递到内函数里面。

def mylog(type):
    def logging(func):
        def inner(*args, **kwargs):
            if type == 'debug':
                print('[DEBUG] logging')
            else:
                print('[INFO] logging')
            rv = func(*args, **kwargs)
            return rv
        return inner
    return logging


@mylog(type='debug')
def add(x, y):
    return x + y


if __name__ == '__main__':
    result = add(1, 2)
    print('result: %d' % result)
    
# 输出：
# [DEBUG] logging
# result: 3

2.4 类装饰器

2.4.1 不带参数的类装饰器

import time


class Timer:
    def __init__(self, func) -> None:
        self.func = func

    def __call__(self, *args, **kwargs):
        start = time.time()
        rv = self.func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        return rv


@Timer
# 等价于 add = Timer(add)
def add(x, y):
    return x + y


if __name__ == '__main__':
    result = add(2, 3)
    print('result: %d' % result)

# 输出：
# time taken 0.000000
# result: 5

上述相当于把一个装饰器变成了一个

Timer

类的对象，然后

add

函数被传入进了

__init__

中，保存为

self.func

，在后面调用

add(2,3)

的时候，实际上相当于调用了

__call__

这个函数，做了一个对象的调用，后面参数2和3就被传入到了

__call__

里面，然后依顺序运行了代码。

2.4.2 带参数的类装饰器

import time


class Timer:
    def __init__(self, prefix) -> None:
        self.prefix = prefix

    def __call__(self, func):
        print('Timer.__call__')

        def wrapper(*args, **kwargs):
            start = time.time()
            ret = func(*args, **kwargs)
            print(f'{self.prefix}:{time.time() - start}')
            return ret

        return wrapper


@Timer(prefix='curr_time:')
# 等价于： add = Timer(prefix = 'curr_time:')(add)
def add(x, y):
    return x + y


if __name__ == '__main__':
    print(add(2, 3))

# 输出：
# Timer.__call__
# curr_time::0.0
# 5

上述把一个装饰器初始化为

Timer

类的对象，然后

prefix

参数被传入进了

__init__

中，之后调用

__call__

函数返回一个闭包。之后就相当于装饰器函数了，用

wrapper

装饰

add

函数。

2.5 @wraps语法糖

因为装饰器实质是就是一个函数，是一个被修饰过函数，他与原来未被修饰的函数是两个不同的函数对象。所以，这个装饰器丢失了原来函数对象的一些属性，比如：

name

，__doc__等属性。

原始函数：

def add(x, y):
    """Add x and y"""
    return x + y


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

# 输出：
# add
# Add x and y

使用装饰器后：

def logging(func):
    def inner(*args, **kwargs):
        """logging.inner"""
        if type == 'debug':
            print('[DEBUG] logging')
        else:
            print('[INFO] logging')
        rv = func(*args, **kwargs)
        return rv
    return inner


@logging
def add(x, y):
    """Add x and y"""
    return x + y


if __name__ == '__main__':
    print(add.__name__)
    print(add.__doc__)
    
# 输出：
# inner
# logging.inner

使用wraps语法糖可以保留原来函数的属性，

@wraps(func)

是个

带参数的类装饰器

。

from functools import wraps


def logging(func):
    @wraps(func)
    def inner(*args, **kwargs):
        """logging.inner"""
        if type == 'debug':
            print('[DEBUG] logging')
        else:
            print('[INFO] logging')
        rv = func(*args, **kwargs)
        return rv
    return inner


@logging
def add(x, y):
    """Add x and y"""
    return x + y


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

# 输出:
# add
# Add x and y