最近在折腾 Telegram Bot 时使用到了 Python 装饰器,顺便记录一下。

建议阅读本文章的 Jupyter Notebook 版本,此版本有附有代码运行结果。

什么是装饰器

装饰器是一个可调用的对象,参数为另一个被装饰的函数,返回一个可调用的对象或函数。

我们先看看 Python 装饰器是干什么的,首先我们先定义一个这样的函数:

def deco(func):
    def inner():
	    print('running %s' % func.__name__)
	    return func()
    return inner

这里的 deco 函数其实就是一个装饰器,将函数作为参数传入,就能使用该装饰器。以下两段代码是等价的,输出的结果一样。

def func1():
    print('running func1()')

func1 = deco(func1)
func1()
@deco
def func2():
	print('running func2()')

func2()

从上面哪个例子可以发现,Python 装饰器其实就是语法糖。所以使用装饰器的时候只要在被装饰函数上面一行使用 @ 符号 + 装饰器即可,无需写一大段代码。

实现一个简单的装饰器

了解了什么是装饰器后,我们来实现一个简单的装饰器。该装饰器可以打印被装饰的函数的运行时间、传入参数和运行结果。

import time

def timethis(func):
    def wrapper(*args):
        start = time.time()
        result = func(*args)
        end = time.time()
        args_str = ', '.join(repr(arg) for arg in args)
        print('[%0.8fs] %s(%s) -> %r' %(end-start, func.__name__, args_str, result))
        return result
    return wrapper

这里给出两个使用该装饰器的示例:

保留函数元信息

使用装饰器时,原函数的元信息都会丢失,这个时候我们就需要使用 functools 库中的 @wraps 装饰器来注解底层包装函数。建议在定义自己的装饰器的时候都加上这个装饰器。

比如,之前我们定义的装饰器就要修改为这样:

import time
from functools import wraps

def timethis2(func):
    @wraps(func)
    def wrapper(*args):
        start = time.time()
        result = func(*args)
        end = time.time()
        args_str = ', '.join(repr(arg) for arg in args)
        print('[%0.8fs] %s(%s) -> %r' %(end-start, func.__name__, args_str, result))
        return result
    return wrapper
# 使用 @wraps 的装饰器 
@timethis2
def accumulate2(n):
    count = 0
    for i in range(n):
        count = count + i
    return count

accumulate2.__name__, accumulate.__name__

从上面的输出可以看出使用了 @wraps 装饰器定义的装饰器保留了原函数的元信息。

解除装饰器

有时候需要使用未被包装的函数,可以使用__wrapped__属性(只有使用了@wraps定义的装饰器才有这个属性)

accumulate2.__wrapped__(6)

未使用 @wrapped 定义的装饰器所装饰的函数没有该属性。

带参数的装饰器

如果想给装饰器添加参数,可以参考下面的例子。

这里我们还是定义一个打印运行时间的装饰器,这里我们定义的时候使用了一个参数 display_result, 当该参数为 False 时不打印函数运行结果。

def timethis3(display_result=True):
    def decorate(func):
        @wraps(func)
        def wrapper(*args):
            start = time.time()
            result = func(*args)
            end = time.time()
            args_str = ', '.join(repr(arg) for arg in args)
            if(display_result):
                print('[%0.8fs] %s(%s) -> %r' %(end-start, func.__name__, args_str, result))
            else:
                print('[%0.8fs] %s(%s)' %(end-start, func.__name__, args_str))
            
            return result
        return wrapper
    return decorate

在定义被装饰的函数的时候,可以在装饰器后面添加参数。

@timethis3(display_result=False)
def accumulate3(n):
    count = 0
    for i in range(n):
        count = count + i
    return count

accumulate3(100)
@timethis3(display_result=True)
def accumulate4(n):
    count = 0
    for i in range(n):
        count = count + i
    return count

accumulate4(100)

装饰器叠加

有时候需要使用到多个装饰器,这个时候该如何定义呢? 直接按顺序将装饰器放在函数上面就行了。

@d1
@d2
def f():
    pass

等价于以下代码:

def f():
    pass

f = d1(d2(f))