Python 装饰器原理

2023-02-06 18:03:10

装饰器是 Python 编程中常用的一个功能，可以将通用的逻辑抽象成装饰器，通过装饰器语法应用到不同的目标上，达到增强或修改目标逻辑的目的。

先来看一个简单的例子

# 打印耗时的装饰器
def log(f):
    def inner(*args, **kw):
        start = time.perf_counter()
        result = f(*args, **kw)
        cost = time.perf_counter() - start
        print(f"cost {cost:.1} seconds")
        return result
    return inner

# 通过 @ 语法将装饰器应用到目标对象上
@log
def download(url):
    requests.get(url)

download("https://www.baidu.com")
# 输出如下
# cost 0.1 seconds

被装饰的目标是 download 函数，用于下载文件，通过 @ 语法被打印耗时的装饰器 log 修饰，download 在被调用的时候就会打印下载文件的耗时。那么打印耗时的逻辑是怎么嵌入到 download 函数中的呢，这里的关键就是需要知道 Python 解释器是怎么处理 @ 语法的。不妨用 Python 提供的 dis 模块反汇编一下 download 的定义

@log
def download(url):
    requests.get(url)

反汇编之后得到的字节码如下（不同 Python 版本反汇编的结果略有差异）

2            LOAD_NAME                 (log)

3           2 LOAD_CONST                (<code object download at 0x7f7872e363a0, file "<dis>", line 2>)
            4 LOAD_CONST               1 ('download')
            6 MAKE_FUNCTION            
            8 CALL_FUNCTION            1
           10 STORE_NAME               1 (download)
           12 LOAD_CONST               2 (None)
           14 RETURN_VALUE

# download 函数本来的函数体
Disassembly of <code object download at 0x7f7872e363a0, file "<dis>", line 2>:
4            LOAD_GLOBAL               (requests)
            2 LOAD_METHOD              1 (get)
            4 LOAD_FAST                 (url)
            6 CALL_METHOD              1
            8 POP_TOP
           10 LOAD_CONST                (None)
           12 RETURN_VALUE

可以看到，一开始 MAKE_FUNCTION 指令通过编译好的 download 函数体对应的 code 对象构造出一个新的函数对象存在 download 变量中，然后通过 CALL_FUNCTION 指令调用 log 函数（参数是 download 变量），后再通过 STORE_NAME 指令将 log 函数的返回结果存到 download 变量中。所以从语义上说，用 log 对 download 进行装饰的效果等价于下面的伪代码

download = make_function(code) # 构造初的 download 函数对象
download = log(download) # 调用 log 函数将返回值赋给 download 变量

log 函数的返回是 inner 函数，也就是说 download 在被修饰之后指向的是 inner 函数，而 inner 函数中的 f 变量又指向了 download 在被修饰之前所指向的函数对象，所以后在调用 download 的时候除了有原来的下载功能之外，还打印了耗时。

装饰前后 download 变量的变化如下图所示

在上面的例子中我们称 log 为装饰器，因为 log 与 @ 同一行且是 @ 后面的内容，这是 Python 语法规定的。log 是一个函数，且定义了一个参数。但是装饰器本身不一定非得是函数，只要是值为可以通过传递一个位置参数进行调用的可调用对象的表达式即可。也就是说

装饰器就是一个表达式
表达式的值是一个可调用对象
这个可调用对象可以通过位置传参的方式传递一个参数进行调用

简单来说，装饰器的本质是可以通过形如 A(f) 的方式进行调用的可调用对象（A 是 @ 后面整个表达式的值）。

可调用对象可以认为是像函数一样可以在后面加括号和参数进行调用的对象，函数本身就是一种可调用对象，还有类对象以及定义了 __call__ 函数的类实例对象等都是可调用对象。

下面的函数 wrapper 就是一个满足条件的装饰器，hello 开始是个空函数，被 wrapper 装饰之后就变成了字符串 "zikcheng"。当然这个装饰器没什么用，只是为了说明定义一个满足语法的装饰器很简单，不一定非要定义内部函数。

def wrapper(a=1, b="hello"):
    return "zikcheng"

@wrapper
def hello():
    pass

print(hello, type(hello))
# 输出
# zikcheng <class 'str'>

再回到 log 与 download 的例子，如果我们想在打印耗时的同时打印负载类型（比如这里是 IO，对于计算密集型的负载类型是 CPU），可以做如下修改

def log_outer(load_type):
    def log(f):
        def inner(*args, **kw):
            start = time.perf_counter()
            result = f(*args, **kw)
            cost = time.perf_counter() - start
            print(f"{load_type} cost {cost:.1} seconds")
            return result
        return inner
    return log

@log_outer("IO")
def download(url):
    requests.get(url)

download("https://www.baidu.com")
# 输出如下
# IO cost 0.1 seconds

这个例子中的装饰器表达式为 log_outer("IO")，因为 @ 符号后面的内容是 log_outer("IO") 整体，而不是 log_outer 。log_outer("IO") 的返回值为 log 函数，因此修饰 download 的装饰器还是 log 函数，只不过这次 log 函数是一个闭包，引用了外部函数的 load_type 变量用于打印负载类型。本质上还是 download = log_outer("IO")(download)。

如果要打印负载类型，除了上面这种方式，还可以通过定义一个类来实现

class LoadLog:
    def __init__(self, load_type):
        self.load_type = load_type

    def __call__(self, f):
        def inner(*args, **kw):
            start = time.perf_counter()
            result = f(*args, **kw)
            cost = time.perf_counter() - start
            print(f"{self.load_type} cost {cost:.1} seconds")
            return result
        return inner

@LoadLog("IO")
def download(url):
    requests.get(url)

download("https://www.baidu.com")
# 输出如下
# IO cost 0.1 seconds