Python 的装饰器简单入门

Python 的装饰器简单入门

一、简单介绍

定义：装饰器是可调用的对象，其参数是另一个函数（被装饰的函数）。装饰器可能会处理被装饰的函数，然后把它返回，或者将其替换成另一个函数或可调用对象。

例如：

@decorate
def target():
    print('running target()')

等价于：

def target():
    print('running target()')

target = decorate(target)

这里 decorate 的返回值必须是一个可调用（callable）对象。

根据上面的定义，可以确定了一个基本的装饰器结构：参数为一个函数，返回值是一个函数（它们可以为同一个，虽然这样没有意义）。

所以，我们可以定义更简单的装饰器：

def useless(f):
    """
    一个最简单的装饰器，参数为一个函数，返回值为一个函数。
    """
    return f

@useless
def speak():
    print("woof, woof, woof")

正如其函数名一样，这是一个无用的装饰器，它什么也没有做，不过这个例子很好的体现了装饰器的定义：参数为一个函数，返回值是另一个函数。

二、几个简单的示例

我们接下来看几个简单的装饰器示例。

1.模拟缓存耗时调用结果的装饰器

import time

INFO_DICT = {
    "key_01": "value_01",
    "key_02": "value_02",
    "key_03": "value_03"
}


def cache(f):
    cache_res = {}
    def wrap(key):
        """
        如果存在就立刻返回，否则执行查询
        """
        start = time.time()
        res = cache_res.get(key)  
        if not res:
            res = f(key)
            cache_res[key] = res
        
        print(f"call func consumes time: {(time.time()-start):.2f} s")
        return res

    return wrap


@cache
def get(key: str) -> str:
    time.sleep(2)
    return INFO_DICT.get(key, "no such key!")


if __name__ == "__main__":
    start = time.time()
    keys = ['key_01', 'key_02', 'key_03', 'key_04', 'key_01', 'key_02', 'key_03']
    for key in keys:
        print(f"{key} => {get(key)}")

输出结果：

call func consumes time: 2.00 s
key_01 => value_01
call func consumes time: 2.00 s
key_02 => value_02
call func consumes time: 2.01 s
key_03 => value_03
call func consumes time: 2.01 s
key_04 => no such key!
call func consumes time: 0.00 s
key_01 => value_01
call func consumes time: 0.00 s
key_02 => value_02
call func consumes time: 0.00 s
key_03 => value_03

2. 模拟日志记录的装饰器

def logger(f):

    def warp(*args, **kwargs):
        # 在调用 f 之前，记录一些信息（这里使用print模拟这个过程）
        print('info: record something.')
        return f(*args, **kwargs)

    return warp

@cache
@logger
def get(key: str) -> str:
    time.sleep(2)
    return INFO_DICT.get(key, "no such key!")

运行结果：

info: record something.
call func consumes time: 2.00 s
key_01 => value_01
info: record something.
call func consumes time: 2.01 s
key_02 => value_02
info: record something.
call func consumes time: 2.01 s
key_03 => value_03
info: record something.
call func consumes time: 2.00 s
key_04 => no such key!
call func consumes time: 0.00 s
key_01 => value_01
call func consumes time: 0.00 s
key_02 => value_02
call func consumes time: 0.00 s
key_03 => value_03

注意看，只有缓存未命中才有日志记录，在缓存命中时，因为根本没有调用 logger，所以就没有日志的输出。如果希望无论缓存是否命中，都记录调用的结果，那么解决办法也很简单，把装饰器的顺序调换一下：

@logger
@cache
def get(key: str) -> str:
    time.sleep(2)
    return INFO_DICT.get(key, "no such key!")

运行结果：

info: record something.
call func consumes time: 2.01 s
key_01 => value_01
info: record something.
call func consumes time: 2.01 s
key_02 => value_02
info: record something.
call func consumes time: 2.01 s
key_03 => value_03
info: record something.
call func consumes time: 2.01 s
key_04 => no such key!
info: record something.
call func consumes time: 0.00 s
key_01 => value_01
info: record something.
call func consumes time: 0.00 s
key_02 => value_02
info: record something.
call func consumes time: 0.00 s
key_03 => value_03

装饰器的顺序就是函数调用的区别：

"""
@cache
@logger
def get(key: str): -> str
    pass
"""
cache(loger(get))

"""
@logger
@cache
def get(key: str): -> str
    pass
"""
log(cache(get))

3. 自动关闭资源的装饰器

def auto_close(f):
    def warp(file=None, *args, **kwargs):
        res = f(file, *args, **kwargs)
        # 调用函数之后，如果file未关闭，则关闭它，防止资源泄露
        # 这里没有做异常处理，未考虑发生异常情况下的资源回收！
        if file:
            print("未关闭文件，自动关闭文件")
            file.close()
        return res
    return warp

@auto_close
def read(file=None)->List[str]:
    return file.readlines()

@auto_close
def write(file=None, line:str=None)->None:
    file.write(line)

    
if __name__ == "__main__":

    dir_path = r'D:/test_case/test'
    # 打开文件进行读写，但是不关闭，让装饰来关闭
    # 对于打开文件，应该使用 with 语句来自动管理，这里仅仅是为了模拟才这样做的
    for i in range(3):
        file_path = os.path.join(dir_path, f"test_{i}.txt")
        input_file = open(file_path, 'a+', encoding="utf-8")
        write(file=input_file, line="To be or not to be, that's the question.")

        input_file = open(file_path, 'a+', encoding="utf-8")
        read(file=input_file)

运行结果：

未关闭文件，自动关闭文件
未关闭文件，自动关闭文件
未关闭文件，自动关闭文件
未关闭文件，自动关闭文件
未关闭文件，自动关闭文件
未关闭文件，自动关闭文件

4.对于可调用对象的简单理解

装饰器的返回值要求是可调用的，所以它本身并没有限制函数嵌套的层级，这意味着多层嵌套也是可能的，让我们来尝试一下吧。

def wrap(greeting:str):
    print(f"wrap => {greeting}")
    def wrap01(greeting:str):
        print(f"wrap01 => {greeting}")
        def wrap02(greeting:str):
            print(f"wrap02 => {greeting}")
            def wrap03(greeting:str):
                print(f"wrap03 => {greeting}")
                def wrap04(f):
                    return f
                return wrap04
            return wrap03
        return wrap02
    return wrap01

wrap_func = wrap("I am wrap func")("I am wrap01 func")("I am wrap02 func")("I am wrap03 func")

@wrap_func
def say_it(greeting:str):
    print(greeting)


if __name__ == "__main__":
    say_it("Look on my works, ye mighty, and despair.")

这里如果写成下面这样，则是错误的:

@wrap("I am wrap func")("I am wrap01 func")("I am wrap02 func")("I am wrap03 func")
def say_it(greeting:str):
    print(greeting)

因为 Python 的语法并不支持这样的写法，尽管上面这个装饰器的返回值是可调用对象，关于这一点查看 Python 的装饰器语法就能明白了：

decorator ::= "@" dotted_name ["(" [argument_list [","]] ")"] NEWLINE

可见，装饰器语法只支持一对可选的圆括号，也就是说最多是支持 直接 的三层嵌套的函数（装饰器）。 不过既然如此，那我就走一个间接的方式来绕过这种方式，所以就产生了上面那种写法。

运行结果：

wrap => I am wrap func
wrap01 => I am wrap01 func
wrap02 => I am wrap02 func
wrap03 => I am wrap03 func
Look on my works, ye mighty, and despair.

这里有一点参数化装饰器的意思，不过这里主要是演示多层嵌套的作用。实际上，这些完全是多余的， 因为普通的装饰器是两层嵌套的函数，如果需要参数化装饰器，则可以再嵌套一层，也就是三层嵌套的函数。 不过，矫枉必过正，为了更直观的理解它，我就设计了这个五层的嵌套函数。