司开星的博客

理解Python函数装饰器

发表于  |    |   分类于 ，   |   阅读次数

Python中的函数装饰器是一种能在不修改函数的前提下给函数添加额外功能的写法。

什么是装饰器

一个函数修改需求

假设已经存在一个函数:

1
2
3
4
import random

def func_a():
    print("I'm func_a, get a random number %s" % random.random())

现在想给函数加一句显示 ，最简单的写法是直接在原函数里加上新代码 :

1
2
3
4
5
import random

def func_a():
    print("I'm new code")
    print("I'm func_a, get a random number %s" % random.random())

不修改原函数的写法

某些情况下我们不好直接修改原函数的代码（实际场景可能是这个函数是别人写的，或者函数逻辑很复杂，不好直接修改，或者同时存在很多类似的函数需要加这个功能），这时候可以再写一个函数，先实现要添加的功能，再运行需要修改的函数并返回结果:

1
2
3
4
5
6
7
8
9
10
import random

def func_a():
    print("I'm func_a, get a random number %s" % random.random())
    
def new_func_a():
    print("I'm new code")
    func_a()

new_func_a()

输出：

1
2
>>> I'm new code
>>> I'm func_a, get a random number 0.582413545375

上面的写法在只有一个函数的情况下可以，但是如果大量函数都需要添加相同功能， 那就要写多个包含print("I'm new func")的新函数，这样写代码的效率太低。

那如果把要修改的函数作为新函数的参数，让新函数来运行新代码呢：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# -*- coding: utf-8 -*-
import random

def func_a():
    print("I'm func_a, get a random number %s" % random.random())
    
def func_b():
    print("I'm func_b, get a random number %s" % random.random())
    
def add_new_code(func): # 目前仅考虑无参数函数的情况
    print("I'm new code")
    func()

new_func_a = add_new_code(func_a)
new_func_b = add_new_code(func_b)

new_func_a()
new_func_b()

乍看符合需求，但实际运行却出了异常：

1
2
3
4
5
6
7
8
I'm new code
I'm func_a, get a random number 0.664406488687
I'm new code
I'm func_b, get a random number 0.66229251997
Traceback (most recent call last):
  File "Untitled 24.py", line 17, in <module>
    new_func_a()
TypeError: 'NoneType' object is not callable

在调用new_func_a()时抛出异常，提示new_func_a是NoneType。原来在调用add_new_code(func_a)时两行print代码已经运行完了，赋给new_func_a的值其实是add_new_code的返回值，没有写返回值的情况下返回值为None。

问题找到了，应该让add_new_code返回一个函数（而不是函数运行结果）：

1
2
3
4
5
6
7
8
def add_new_code(func):
    print("I'm new code")
    '''
    这里是func而不是func()
    func()是调用并返回运行结果，即 new_func_a 获取的仍然只是单次运行的值，无法作为函数调用
    返回func则是将原函数赋给 new_func_a，这样new_func_a 才是一个函数
    '''
    return func

再实际运行一下得到结果：

1
2
3
4
I'm new func
I'm new func
I'm func_a, get a random number 0.489578986299
I'm func_b, get a random number 0.878320346664

结果顺序好像不太对，为什么呢？

问题与上面的问题本质相同，print("I'm new code")是在调用add_new_code时运行的，而不是在new_func_a运行，即赋给new_func_a的值其实还是原来的func_a，并没有添加上新代码。

现在我们来梳理一下add_new_code要实现的功能：函数作为参数，返回一个新函数，新函数需要运行新代码及原函数两部分。

我们先来写新函数（即应该被add_new_code返回的函数）：

1
2
3
def new_func(): # 不能带参数
    print("I'm new code")
    return func()

注意这里new_func不能带参数，因为新函数应与原函数接收的参数保持一致。

要返回的函数已经有了，现在要解决的问题就是：这里的func从哪来？

实现装饰器

解决办法是Python中的闭包（即内部函数使用外部函数作用域的值）。写法如下：

1
2
3
4
5
6
7
8
9
10
11
def add_new_code(func):
    def new_func():
        print("I'm new code")
        return func()
    return new_func

new_func_a = add_new_code(func_a)
new_func_b = add_new_code(func_b)

new_func_a()
new_func_b()

运行结果：

1
2
3
4
I'm new code
I'm func_a, get a random number 0.406357569132
I'm new code
I'm func_b, get a random number 0.414812846958

结果正是我们需要的。这种写法就是装饰器。

实际项目中一般不愿意为一个小改动就改一次函数名，特别是不影响原函数实际功能的改动。上面的代码可以优化一下：

func_a = add_new_code(func_a)

即把原函数名指向装饰后的新函数，这样调用原函数的代码也不需要修改了。

Python为装饰器提供了更方便的写法（语法糖）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# -*- coding: utf-8 -*-
import random

def add_new_code(func):
    def new_func():
        print("I'm new code")
        return func()
    return new_func

# 在原函数前加上@add_new_code 相当于func_a = add_new_code(func_a)
@add_new_code  
def func_a():
    print("I'm func_a, get a random number %s" % random.random())

@add_new_code
def func_b():
    print("I'm func_b, get a random number %s" % random.random())
    
func_a()
func_b()

Python 同时还为装饰器提供了保留原函数元信息（比如func_a.__name__获取的函数名等）的方式：

1
2
3
4
5
6
7
8
from functools import wraps

def add_new_code(func):
    @wraps(func)  # 在内部函数前加上@wraps装饰器
    def new_func():
        print("I'm new code")
        return func()
    return new_func

这就是标准的装饰器写法了。

不同需求下的装饰器

原函数带参数

前面考虑的都是原函数不带参数的情况，如果带参数则要修改一下装饰器函数：

1
2
3
4
5
6
7
8
from functools import wraps

def add_new_code(func):
    @wraps(func)  # 在内部函数前加上@wraps装饰器
    def new_func(*args, **kwargs):
        print("I'm new code")
        return func(*args, **kwargs)
    return new_func

*args, **kwargs确保原函数不管是什么参数此装饰器都适用。

带参数的装饰器

有时候装饰器函数需要根据情况不同运行不同代码，每种情况都单独写一个装饰器会出现太多重复代码（Don’t repeat yourself）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from functools import wraps

def add_new_code_one(func):
    @wraps(func)
    def new_func(*args, **kwargs):
        print("I'm new code")
        print("I'm new code one")
        return func(*args, **kwargs)
    return new_func
    
def add_new_code_two(func):
    @wraps(func)
    def new_func(*args, **kwargs):
        print("I'm new code")
        print("I'm new code two")
        return func(*args, **kwargs)
    return new_func

要在一个装饰器函数里根据情况运行不同代码就要把不同情况作为一个参数传递给装饰器函数。但是add_new_code已经接收了func作为参数，新参数要怎么传呢？

考虑装饰过程的实际代码：

func_a = add_new_code(func_a)

如果改成这样：

func_a = add_new_code(arg)(func_a)

即add_new_code(arg)返回的才是装饰后的函数。

按照这种思路，装饰器最外层的函数要接收的就应该是不同情况这个参数：

1
2
3
4
5
6
7
8
9
10
11
from functools import wraps

def add_new_code(arg):
    def decorate(func):
        @wraps(func)
        def new_func(*args, **kwargs):
            print("I'm new code")
            print("Get a arg: %s" % arg)
            return func(*args, **kwargs)
        return new_func
    return decorate

即在外部函数外再套一层外部函数，主要是为了让最里层的函数能同时使用func和arg两个参数。

装饰原函数的语法糖写法如下:

1
2
3
@add_new_code('wrapper for func_a') 
def func_a():
    print("I'm func_a, get a random number %s" % random.random())

@之后的部分仍然是一个把被装饰函数作为参数的函数。

那为什么不用另一种写法：

func_a = add_new_code(func_a)(arg)

按这种写法装饰器函数应该这么写：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from functools import wraps

def add_new_code(func):
    def decorate(arg):
        @wraps(func)
        def new_func(*args, **kwargs):
            print("I'm new code")
            print("Get a arg: %s" % arg)
            return func(*args, **kwargs)
        return new_func
    return decorate

def func_a():
    print("I'm func_a, get a random number %s" % random.random())

func_a = add_new_code(func_a)("wrapper for func_a")
func_a()

运行结果与前面的结果相同。

虽然运行结果相同但这种写法没有@语法糖的支持，只能手写add_new_code(func_a)("wrapper for func_a")，并且与标准写法相比没有优势，此处写出来只为让读者对比以理解装饰器的本质。

将类定义为装饰器

装饰器可以装饰的实际上是可调用对象（函数就是最常见的可调用对象），对于类来说可调用对象需要实现__call__方法。类作为装饰器的写法如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import random
from functools import wraps

class Profiled:
    def __init__(self, func):
        self.func = func
        self.call_count = 0 # 装饰后的函数运行次数
        self = wraps(func)(self) # wraps装饰器不使用@的写法

    def __call__(self, *args, **kwargs):
        self.call_count += 1
        print("I'm new code")
        print("new func running count: %s" % self.call_count)
        return self.func(*args, **kwargs)

@Profiled    
def func_a():
    print("I'm func_a, get a random number %s" % random.random())
    
func_a()
func_a()
print(func_a.__name__) # 检查元信息

输出:

1
2
3
4
5
6
7
I'm new code
new func running count: 1
I'm func_a, get a random number 0.710760549208
I'm new code
new func running count: 2
I'm func_a, get a random number 0.507024218664
func_a

由于类实例的初始化和运行过程分离，可以实现这种运行计数功能。

不过类装饰器常见的用法是与描述器结合，关于描述器可参考本人的另一篇文章。

示例

下面两个来自《Python Cookbook》，都是实际项目中常见的需求：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import time
from functools import wraps
def timethis(func):
    '''
    输出函数运行时间的装饰器
    '''
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end-start)
        return result
    return wrapper
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from functools import wraps
import logging

def logged(level, name=None, message=None):
    """
    为函数添加日志。
    level参数是日志等级
    message参数是日志记录的额外信息
    """
    def decorate(func):
        logname = name if name else func.__module__
        log = logging.getLogger(logname)
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            log.log(level, logmsg)
            return func(*args, **kwargs)
        return wrapper
    return decorate