Python装饰器

发表于 | 分类于 |

最近面网易互娱的时候，面试官看到简历上写了Python，于是问我Python装饰器会不会。Ummmm好像还真不会，最后面试官说我Python没有掌握得很好。想来也是，自己使用Python无非就是用来写各种算法题和各种调库，很少去学Python本身的一些东西（比方说修饰器，多线程，变量作用域等等）。讲道理Python在这么多编程语言里面应该算是很好学的了，感觉还是该花时间好好学一下，让它变成自己真正掌握的东西。

那我就先从装饰器开始。

为什么需要装饰器

这里用一个很经典的例子。假设我有一个简单的函数：

1
2
3
def f(a, b):
    c = a + b
    print("The result is : %d" % (c))

现在我们想要统计函数的运算时间，那么我们正常的思路是对函数进行修改，加上一个计时器的功能。

1
2
3
4
5
6
7
8
9
import time
def f(a, b):
    start = time.time()

    c = a + b
    print("The result is : %d" % (c))

    end = time.time()
    print("Total time is : %d" % (end - start))

但很多时候我们是不能直接修改核心函数本身的，因为这便混淆了函数原本的功能，失去了封装性。另外，如果我们希望为函数包装上多种不同的功能，比方说多种精度的计时器，那我们需要频繁地修改函数本身，而且不方便测试。

那么我们可以尝试用另一种方法，新定义一个函数，将原函数作为输入，然后在新函数里实现计时器的功能。

1
2
3
4
5
6
7
8
9
10
11
12
import time
def f(a, b):
    c = a + b
    print("The result is : %d" % (c))

def timer(func, a, b):
    start = time.time()

    func(a, b)

    end = time.time()
    print("Total time is : %d" % (end - start))

但这也会带来其它的问题。如果你在代码的不同角落都调用了f(a, b)，那你需要在这些所有地方都用新的timer替换f。

回过头来想，这个做法实际上就是用一个新的函数来包装原函数本身，然后在调用原函数的基础上增添新的功能，于是便有了装饰器。

装饰器怎么搞

装饰器的用法其实很简单，我们只需要定义一个新的函数，这个函数接受原函数作为参数，然后返回一个与原函数函数签名相同（即参数和返回值相同）的函数，最后用 @deco 关键字进行修饰即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
def deco(f):
    def wrapper(a, b):
        start = time.time()
        f(a, b)
        end = time.time()
        print("Total time is : %d" % (end - start))
    return wrapper

@ deco
def func(a, b):
    c = a + b
    print("The result is : %d" % (c))
func(1, 2)

通过这样的方式，我们便将原函数替换为新的函数（添加了新的功能），并且不需要在代码本身及每个调用函数的地方都进行修改。

更泛化的装饰器

从上面的装饰器我们能看到一个问题，我们这个用作计时器的装饰器实际上功能有限，只能用来装饰带两个参数没有返回值的函数。那如果我们想要用一个计时器装饰器来装饰所有的函数，该怎么做呢？

这里我们可以用到python中的可变参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def deco2(f):
    def wrapper(*args, **kwargs):
        print("Deco 2")
        start = time.time()
        f(*args, **kwargs)  # 这里的传递方式等价于 f(args[0], args[1])，因为加了个星号
        end = time.time()
        print("Total time : %d" % (end - start))
    return wrapper

@ deco2
def func(a, b):
    c = a + b
    print("Result is : %d" % (c))
func(1, 2)

这样，我们就可以通过这个装饰器来修饰几乎所有的函数了。

多重装饰

如果我们想要为原函数添加多个不同的功能（比方说，计时并打印当前线程的id），该怎么办？很简单，因为修饰器是可以嵌套的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def deco(f):
    def wrapper(a, b):
        print("Deco 1")
        start = time.time()
        f(a, b)
        end = time.time()
        print("Total time : %d" % (end - start))
    return wrapper

def deco2(f):
    def wrapper(*args, **kwargs):
        print("Deco 2")
        start = time.time()
        c = f(*args, **kwargs)
        end = time.time()
        print("Total time : %d" % (end - start))
        
    return wrapper

@ deco
@ deco2
def func(a, b):
    c = a + b
    print("Result is : %d" % (c))
func(1, 2)

输出如下：

1
2
3
4
5
Deco 1
Deco 2
Result is : 3
Total time : 0
Total time : 0

注意这里嵌套的顺序为：deco(deco2(func))

参考：python装饰器的详细解析