文章目录

一文掌握Python装饰器

一文掌握Python装饰器

1. 什么是装饰器

举个例子，我们在写Python类的时候，为了将成员设置为只读类型的，我们可以用@property装饰器去装饰一个与成员变量同名的函数，如下：

class Student:
    def __init__(self):
        self._name = "Jack"

    @property
    def name(self):
        return self._name

s = Student()
print(s.name)

装饰器在Python中是很常见的功能，像Flask、Django这种WEB框架中大量被使用，在MMCV开源项目中，也用到了装饰器来实现注册功能。那么装饰器到底是个什么东西？底层又是怎么实现的呢？

所谓装饰器(decorator)，是python中的一个可调用对象(callable),我们可以简单地将它理解为一个函数，只不过这个函数是包装了另一个函数的函数。我们不妨用代码来解释一下它的形式：

# Case 1
@decorate
def target():
  print("Running target()")
  
# Case 2 
def target():
  print("Running target()")
target=decorate(target)

以上两种写法的效果完全等价，由于Python中秉承“一切皆对象”的原则，函数也是一个对象，装饰器的功能就是把原来的函数包装成一个新的函数对象返回，这个新的函数对象可以在原函数功能的基础之上，扩展一些其他的功能。

2. 装饰器基础知识

为了验证我们上文说的装饰器会返回一个新的函数对象，我们写一个简单的代码验证一下：

def deco(func):
    def inner():
        print("Running inner()")
    return inner

@deco
def target():
    print("Running target()")

target()
print(target)

# output
Running inner()
<function deco.<locals>.inner at 0x7ff87e50fc80>

可以看到，返回的新的函数对象其实就对inner()函数的引用，而且我们这个装饰器的直接把原函数target()的功能屏蔽掉了。这里有个小问题，Python会在什么时候执行装饰器呢？我们也不妨用一个例子来测试一下。

registry = []

def register(func):
    print("Running register <%s>" % func)
    registry.append(func)
    return func

@register
def func1():
    print("Running func1()")

@register
def func2():
    print("Running func2()")

def func3():
    print("Running func3()")

if __name__ == '__main__':
    print("Running main()")
    print("Registry --> ", registry)
    func1()
    func2()
    func3()

我们这里模拟一个函数注册中心的功能，实现了一个名为register的装饰器，运行以上代码，我们会得到一下输出：

Running register <<function func1 at 0x7f9f37b0fbf8>>
Running register <<function func2 at 0x7f9f37b0fc80>>
Running main()
Registry -->  [<function func1 at 0x7f9f37b0fbf8>, <function func2 at 0x7f9f37b0fc80>]
Running func1()
Running func2()
Running func3()

这说明了一个问题，在上述脚本程序启动之前，装饰器就已经开始包装函数，所以，我们可以得到一个结论——装饰器在被装饰得函数定义之后立即执行！也就是说，在写了@XXX之后，装饰器就开始装饰了，而不是等到运行时才执行。

2.1 变量作用域

作用域，无非可分为两类——全局变量和局部变量，这是在任何编程语言中都会存在的概念。通常来说，在Python函数中，函数参数和函数内部赋值的变量，我们都称为局部变量，而函数外面的变量我们都称为全局变量。当然，任何变量在被访问之前，都必须有定义。例如，

def f1(a):
  print(a)
  print(b)

f1()
# output
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in f1
NameError: name 'b' is not defined

显然，变量b应该是个全局变量，但是我们并没有预先声明，那自然就报错了，这个例子很好理解，但下面的这个例子可能会出乎我们的意料。

b=6
def f2(a):
  print(a)
  print(b)
	b=9

f2(3) 
# output
3
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in f2
UnboundLocalError: local variable 'b' referenced before assignment

是不是有点惊讶？这里Python解释器把变量b理解成了局部变量，我们可能应该会将这个b理解成全局变量，并且打印出6把！可事实就是，Python编译函数的定义体时，发现变量b被赋值了，于是就将b解释成局部变量。那你可能会问，如果这里我就是想要使用那个全局变量b，怎么做呢？答案是，用global关键字显式声明！如下：

b=6
def f3(a):
  global b
  print(a)
  print(b)
  b=9

f3(3)
print(b)
# output
3
6
9

在上面的例子中，变量b始终被认为是全局变量，在函数体内修改也会对全局变量的值造成影响。为了更加坚定我们的结论，不放将上述代码反汇编一下，用Python自带的dis包即可完成。

>>> from dis import dis
>>> dis(f1)
  2           0 LOAD_GLOBAL              0 (print)
              2 LOAD_FAST                0 (a)
              4 CALL_FUNCTION            1
              6 POP_TOP

  3           8 LOAD_GLOBAL              0 (print)
             10 LOAD_GLOBAL              1 (b)   # 变量b在f1()中被解释为全局变量
             12 CALL_FUNCTION            1
             14 POP_TOP
             16 LOAD_CONST               0 (None)
             18 RETURN_VALUE

>>> dis(f2)
  2           0 LOAD_GLOBAL              0 (print)
              2 LOAD_FAST                0 (a)
              4 CALL_FUNCTION            1
              6 POP_TOP

  3           8 LOAD_GLOBAL              0 (print)
             10 LOAD_FAST                1 (b)   # 变量b在f1()中被解释为局部变量
             12 CALL_FUNCTION            1
             14 POP_TOP

  4          16 LOAD_CONST               1 (9)
             18 STORE_FAST               1 (b)
             20 LOAD_CONST               0 (None)
             22 RETURN_VALUE

2.2 闭包

闭包(closure),这个词听上去就很专业很高级的样子，其实很简单，它指的就是延伸了作用域的函数，其中包含了，在函数定义体中引用了，可是又在定义体中未曾定义过的非全局变量。好绕口啊，你可以先尝试理解一下这句话。好吧，我猜你肯定没明白，还是通过例子来说明一下。现在有这样一个需求，需要根据不断到来的数据，求他们的平均值，你怎么实现？我们希望你可以定义一个名为avg()的函数，使用方法大概如下：

>>> avg(10)
10.0
>>> avg(11)
10.5
>>> avg(12)
11.0

第一反应可能是这样：

class Averager:
    def __init__(self):
        self.vals = []

    def __call__(self, val):
        self.vals.append(val)
        return sum(self.vals) / len(self.vals)

avg = Averager()
print(avg(10))
print(avg(11))
print(avg(12))

或者，高级一点，可以用一个高阶函数来实现，如下:

def make_averager():
    vals = []  # Notice here

    def averager(val):
        vals.append(val)
        return sum(vals) / len(vals)

    return averager

avg = make_averager()
print(avg(10))
print(avg(11))
print(avg(12))

其实这两个思路都是用一个数组来保存历史值，然后求出平均值来。可是问题是，对于函数make_averager(),变量vals是局部变量没问题，可是调用avg(10)的时候，函数make_averager()已经返回了，局部变量的作用域也一去不复返了，在内部函数averager()中，应该报错才是呀！可事实就是，它运行的好好的，因为对于averager()来说，变量vals是一个自由变量(free variable)，这是一个技术术语，专门指那些未在本地作用域中绑定的变量。

同样，为了坚信我们的结论，我们可以从Python的__code__（编译后的函数定义体）属性中寻找线索，该属性是一个字典，其中包含了key为co_varnames的局部变量元组，以及key为co_freevars的自由变量元组，打印我们分别可以得到输出('val',)和('vals',)。

3. 实现简单的装饰器

我们说，装饰器可以用来装饰原函数，在原函数的基础之上增加一些功能，比如常用的日志Hook、消息通知等，那我们怎么实现一个简单的装饰器呢？其实在第2小节的开头，我们已经实现了一个注册功能的装饰器@register，只不过这个装饰器只是简单地将函数注册到列表中，并且屏蔽了原函数的功能。假设我们现在有这样一个需求，希望在统计一下函数运行的时间，怎么做呢？

笨一点的办法如下：

import time

def func():
  _tic=time.time()
  # ...
  # 函数原来的执行逻辑
  _toc=time.time()
	print(_toc-_tic)

那么换成装饰器怎么实现呢？你可以在心里默念，装饰器接受函数对象作为输入，然后返回一个新的函数对象。开始写代码吧：

# 计时装饰器
def clock(func):
    def inner(*args):
        _tic = time.time()
        res = func(*args)
        elapsed = time.time() - _tic
        func_name = func.__name__
        arg_str = ", ".join([repr(arg) for arg in args])
        print("[%.8fs] %s(%s) --> %r" % (elapsed, func_name, arg_str, res))

        return res

    return inner


@clock
def factorial(n):
    return 1 if n < 2 else n * factorial(n-1)

factorial(1)
print(factorial.__name__)  # 注意这一行的输出
factorial(5)
print()
factorial(10)

[0.00000095s] factorial(1) --> 1
inner
[0.00000095s] factorial(1) --> 1
[0.00000620s] factorial(2) --> 2
[0.00001097s] factorial(3) --> 6
[0.00001502s] factorial(4) --> 24
[0.00001907s] factorial(5) --> 120

[0.00000000s] factorial(1) --> 1
[0.00000429s] factorial(2) --> 2
[0.00000715s] factorial(3) --> 6
[0.00001097s] factorial(4) --> 24
[0.00001407s] factorial(5) --> 120
[0.00001788s] factorial(6) --> 720
[0.00002098s] factorial(7) --> 5040
[0.00002384s] factorial(8) --> 40320
[0.00002789s] factorial(9) --> 362880
[0.00003290s] factorial(10) --> 3628800

这个装饰器已经完成了我们的需求，但是并不是实现的很完美，主要有以下两个问题：

因为inner()函数接受的是*args，暂时无法输入一些形如key-value的关键字参数；
被装饰的函数factorial()的名字变成了inner，这也很好理解，因为装饰器的确返回的是inner()函数对象，也就是说，装饰后的factorial()其实就是inner()函数

以上第一个问题，比较好解决，在inner()函数的定义中加上**kwargs即可；而对于第二个问题，我们需要用Python内置的装饰器@functools.wrap()来包装一下inner()函数，是不是有点套娃的感觉了。。。改写一下上一个版本：

import time
import functools


# 计时装饰器
def clock(func):
    @functools.wraps(func)
    def inner(*args, **kwargs):
        _tic = time.time()
        res = func(*args, **kwargs)
        elapsed = time.time() - _tic
        func_name = func.__name__

        arg_list = []
        if args:
            arg_list.append(", ".join([repr(arg) for arg in args]))
        if kwargs:
            key_values = ["%s=%r" % (k, v) for k, v in sorted(kwargs.items())]
            arg_list.append(", ".join(key_values))
        arg_str = ", ".join(arg_list)
        print("[%.8fs] %s(%s) --> %r" % (elapsed, func_name, arg_str, res))

        return res

    return inner


@clock
def factorial(n, welcome_str=None):
    if welcome_str:
        print(welcome_str)
    return 1 if n < 2 else n * factorial(n - 1)


factorial(1, welcome_str="===== Start from 1 =====")
print(factorial.__name__)
factorial(5, welcome_str="===== Start from 5 =====")
print()
factorial(10, welcome_str="===== Start from 10 =====")

# 输出
===== Start from 1 =====
[0.00003576s] factorial(1, welcome_str='===== Start from 1 =====') --> 1
factorial
===== Start from 5 =====
[0.00000095s] factorial(1) --> 1
[0.00000715s] factorial(2) --> 2
[0.00001097s] factorial(3) --> 6
[0.00001526s] factorial(4) --> 24
[0.00002217s] factorial(5, welcome_str='===== Start from 5 =====') --> 120

===== Start from 10 =====
[0.00000000s] factorial(1) --> 1
[0.00000429s] factorial(2) --> 2
[0.00000906s] factorial(3) --> 6
[0.00001192s] factorial(4) --> 24
[0.00001597s] factorial(5) --> 120
[0.00002003s] factorial(6) --> 720
[0.00002480s] factorial(7) --> 5040
[0.00002885s] factorial(8) --> 40320
[0.00003195s] factorial(9) --> 362880
[0.00003791s] factorial(10, welcome_str='===== Start from 10 =====') --> 3628800

修改过后，我们发现上述提及的两个得以解决，这里需要注意的是，经过@functools.wrap()包装后，装饰器返回的仍然是inner()函数这个对象，只不过这个对象的名字被强制改成了factorial。

细心的你可能会发现，@functools.wrap()这个装饰器居然可以接受一个参数作为输入，而我们自定义的@clock装饰器是无法接受参数作为输入的，这说明我们的装饰器仍然不够完善！

4. 参数化装饰器

怎么让装饰器接受参数呢？一个简单的想法就是，我们写一个装饰器工厂函数，根据入参的不同，构造不同的装饰器，然后再用这个装饰器去装饰原函数。看到这句话肯定又晕了，我们回到之前的@register装饰器的例子。在第2小节中，我们实现的@register装饰器只能注册，而不能注销，现在我们就是想扩展注销的功能。稍作修改，代码如下：

registry = set()

def register(flag=True):
    def decorate(func):
        if flag:  # flag是一个自由变量
            registry.add(func)
            print("Register func <%s>" % func.__name__)
        else:
            registry.discard(func)
            print("Unregister func <%s>" % func.__name__)
        return func

    return decorate


@register(flag=True)
def func1():
    print("Running func1()")


@register()
def func2():
    print("Running func2()")


if __name__ == '__main__':
    print("Running main()")
    print("Registry --> ", registry)
    func1()
    func2()
    register(flag=False)(func2)
    print("Registry --> ", registry)
    
# 输出
Register func <func1>
Register func <func2>
Running main()
Registry -->  {<function func1 at 0x7fda2b1d5c80>, <function func2 at 0x7fda2b1d5d08>}
Running func1()
Running func2()
Unregister func <func2>
Registry -->  {<function func1 at 0x7fda2b1d5c80>}

心里再次默念（多念几遍就会写代码了），装饰器接受函数对象作为输入，返回一个函数对象。好，我们看下现在的装饰器@register，严格来说，它并不是装饰器，而是装饰器工厂，因为它的输入不是函数对象，而是参数flag，它的返回值才是装饰器，也就是我们的@decorate！然后我们再用这个装饰器去装饰函数就行了。噢，这下是真的明白了！

5. 完整的装饰器例子

为了总结本文，我们不妨写一个完整的装饰器，为了说明问题，但又不至于混杂太多复杂逻辑，我们就继续扩展上述的计时器的例子，需要实现以下功能：

不能改变原函数的执行逻辑的前提下，记录函数的执行时间；
用户可以指定保存日志信息的地方，例如stdout、文件等。

开始写代码，不对，是开始默念，参数化装饰器是装饰器工厂，装饰器接受函数对象作为输入，返回新的函数对象，念完再写代码：

import time
import functools


# 计时装饰器
def clock(log_out="stdout"):
    assert log_out in ["stdout", "log.txt"], "目前只标准输出和日志文件输出！"

    def decorator(func):
        if log_out == "stdout":
            logger = print
        else:
            logger = open("log.txt", "w+").write

        @functools.wraps(func)
        def inner(*args, **kwargs):
            _tic = time.time()
            res = func(*args, **kwargs)
            elapsed = time.time() - _tic
            func_name = func.__name__

            arg_list = []
            if args:
                arg_list.append(", ".join([repr(arg) for arg in args]))
            if kwargs:
                key_values = ["%s=%r" % (k, v) for k, v in sorted(kwargs.items())]
                arg_list.append(", ".join(key_values))
            arg_str = ", ".join(arg_list)
            # 记录日志
            logger("[%.8fs] %s(%s) --> %r\n" % (elapsed, func_name, arg_str, res))

            return res

        return inner

    return decorator


@clock(log_out="stdout")
def factorial(n, welcome_str=None):
    if welcome_str:
        print(welcome_str)
    return 1 if n < 2 else n * factorial(n - 1)


@clock(log_out="log.txt")
def fibbo(n, welcome_str=None):
    if welcome_str:
        print(welcome_str)

    return n if n < 2 else fibbo(n - 1) + fibbo(n - 2)


factorial(5, welcome_str="===== Start counting 5! =====")
print(factorial.__name__)

fibbo(5, welcome_str="===== Start counting fibbo(5) =====")
print(fibbo.__name__)

# 输出
===== Start counting 5! =====
[0.00000191s] factorial(1) --> 1
[0.00005794s] factorial(2) --> 2
[0.00006723s] factorial(3) --> 6
[0.00007892s] factorial(4) --> 24
[0.00018406s] factorial(5, welcome_str='===== Start counting 5! =====') --> 120
factorial

===== Start counting fibbo(5) =====
fibbo

连乘的运算日志在stdout，而斐波那契函数的日志在log.txt文件中。

6. 总结

看完本文，你应该可以回答以下的问题：

装饰器是什么？即修饰了原函数的新函数对象。
怎么实现装饰器？装饰器接受函数对象作为输入，返回新的函数对象
怎么让装饰器返回的函数对象和原函数同名？用@functools.warp()装饰器包装装饰器中定义的inner()函数。
怎么实现参数化装饰器？用装饰器工厂来实现！

7. 参考文献

[1] Luciano Remalho, Fluent Python, Function Decorators and Closures, 2015

一文掌握Python装饰器

@decorator

一文掌握Python装饰器

1. 什么是装饰器

2. 装饰器基础知识

2.1 变量作用域

2.2 闭包

3. 实现简单的装饰器

4. 参数化装饰器

5. 完整的装饰器例子

6. 总结

7. 参考文献

FEATURED TAGS

FRIENDS