Python实现简单的装饰器

Python的装饰器是一种嵌套函数，一种闭包。是很好用的工具，利用得好，让代码锦上添花、如虎添翼。本文介绍简单的装饰器的实现，以及三个好用的Python自带装饰器：functools.wraps、functools.lru_cache、functools.singledispatch

一个简单的装饰器，，输出函数的运行时间，代码示例：

from time import sleep, perf_counter as pc 

def clock(func):

    def clocked(*args, **kwargs):
        '''I'm clocked function'''
        t = pc()
        result = func(*args, **kwargs)  # 执行主函数 重要
        elapsed = pc() - t # 主函数执行时间 重要

        # func的属性 重要
        name = func.__name__

        # 格式化参数  打印信息  不重要
        arg_str = ', '.join(repr(arg) for arg in args)
        if not kwargs:  
            print('[%0.8fs] %s(%s) -> %r' % (elapsed, name, arg_str, result))
        else:
            print('[%0.8fs] %s(%s ...) -> %r' % (elapsed, name, arg_str, result))

        return result
        
    return clocked

测试用例：

@clock
def snooze(seconds):
    '''I'm snooze function'''
    sleep(seconds)

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

if __name__ == '__main__':
    # 打印被装饰函数的属性 看有没有被覆盖
    name = snooze.__name__
    doc = snooze.__doc__
    print(name, doc)

    print(snooze(.123))
    print(factorial(6))

输出结果：

clocked I'm clocked function
[0.12802433s] snooze(0.123) -> None
None
[0.00000182s] factorial(1) -> 1
[0.00003298s] factorial(2) -> 2
[0.00005212s] factorial(3) -> 6
[0.00006871s] factorial(4) -> 24
[0.00008502s] factorial(5) -> 120
[0.00010664s] factorial(6) -> 720
720

可以看出装饰器的典型行为：把被装饰的函数替换成新函数，二者接受相同的参数，而且（通常）返回被装饰的函数本该返回的值，同时还会做些额外操作。

怎么保留被装饰的函数原来的属性，比如__name__、__doc__呢？

代码改进：

from time import sleep, perf_counter as pc 
import functools

def clock(func):
    @functools.wraps(func) # 不覆盖被装饰的函数的属性
    def clocked(*args, **kwargs):
        '''I'm clocked function'''
        t = pc()
        result = func(*args, **kwargs)  # 执行主函数 重要
        elapsed = pc() - t # 主函数执行时间 重要

        # func的属性 重要
        name = func.__name__

        # 格式化参数  打印信息  不重要
        arg_str = ', '.join(repr(arg) for arg in args)
        if not kwargs:  
            print('[%0.8fs] %s(%s) -> %r' % (elapsed, name, arg_str, result))
        else:
            print('[%0.8fs] %s(%s ...) -> %r' % (elapsed, name, arg_str, result))

        return result
        
    return clocked

测试输出 snooze.__name__ 发现起作用了，被装饰的函数的属性没有被覆盖。

functools.wraps 只是标准库中拿来即用的装饰器之一。 functools 模块中最让人印象深刻的两个装饰器要属lru_cache 和 singledispatch。

functools.lru_cache是非常实用的装饰器，可以用来缓存，这是一项优化技术，它把耗时的函数的结果保存起来，避免传入相同的参数时重复计算。LRU 三个字母是“Least Recently Used”的缩写，表明缓存不会无限制增长，一段时间不用的缓存条目会被扔掉。

生成第 n 个斐波纳契数这种慢速递归函数特别适合使用 lru_cache，效率会有质的提升：

@clock
def fibonacci(n):
    return n if n < 2 else fibonacci(n-1) + fibonacci(n-2)

if __name__ == '__main__':
    print(fibonacci(6))

输出：

[0.00000060s] fibonacci(1) -> 1
[0.00000076s] fibonacci(0) -> 0
[0.00006251s] fibonacci(2) -> 1
[0.00000043s] fibonacci(1) -> 1
[0.00007643s] fibonacci(3) -> 2
[0.00000039s] fibonacci(1) -> 1
[0.00000039s] fibonacci(0) -> 0
[0.00001284s] fibonacci(2) -> 1
[0.00010172s] fibonacci(4) -> 3
[0.00000036s] fibonacci(1) -> 1
[0.00000036s] fibonacci(0) -> 0
[0.00001278s] fibonacci(2) -> 1
[0.00000035s] fibonacci(1) -> 1
[0.00002500s] fibonacci(3) -> 2
[0.00013924s] fibonacci(5) -> 5
[0.00000033s] fibonacci(1) -> 1
[0.00000036s] fibonacci(0) -> 0
[0.00001243s] fibonacci(2) -> 1
[0.00000040s] fibonacci(1) -> 1
[0.00002480s] fibonacci(3) -> 2
[0.00000033s] fibonacci(1) -> 1
[0.00000057s] fibonacci(0) -> 0
[0.00001910s] fibonacci(2) -> 1
[0.00007221s] fibonacci(4) -> 3
[0.00022999s] fibonacci(6) -> 8
8

可以看到有很多重复计算的，使用functools.lru_cache优化：

@clock
@functools.lru_cache()
def fibonacci(n):
    return n if n < 2 else fibonacci(n-1) + fibonacci(n-2)

if __name__ == '__main__':
    print(fibonacci(6))

输出结果：

[0.00000098s] fibonacci(1) -> 1
[0.00000098s] fibonacci(0) -> 0
[0.00007543s] fibonacci(2) -> 1
[0.00000063s] fibonacci(1) -> 1
[0.00009011s] fibonacci(3) -> 2
[0.00000048s] fibonacci(2) -> 1
[0.00011252s] fibonacci(4) -> 3
[0.00000036s] fibonacci(3) -> 2
[0.00012652s] fibonacci(5) -> 5
[0.00000035s] fibonacci(4) -> 3
[0.00014190s] fibonacci(6) -> 8
8

可以看到functools.lru_cache的缓存效果是挺明显的，而且会随着参数的增大效果变得越明显。效率是质的提升！

注意，必须像常规函数那样调用 lru_cache。这一行中有一对括号：@functools.lru_cache()。这么做的原因是，lru_cache 可以接受配置参数。完整的参数是functools.lru_cache(maxsize=128, typed=False)

maxsize 参数指定存储多少个调用的结果。缓存满了之后，旧的结果会被扔掉，腾出空间。为了得到最佳性能，maxsize 应该设为 2 的幂。typed 参数如果设为 True，把不同参数类型得到的结果分开保存，即把通常认为相等的浮点数和整数参数（如 1 和 1.0）区分开。顺便说一下，因为 lru_cache 使用字典存储结果，而且键根据调用时传入的定位参数和关键字参数创建，所以被 lru_cache 装饰的函数，它的所有参数都必须是可散列的。

单分派泛函数

先看示例：

import html

def htmlize(obj):
    content = html.escape(repr(obj))
    return '<pre>{}</pre>'.format(content)

思考：这个函数适用于任何 Python 类型，但是现在我们想做个扩展，让它使用特别的方式显示某些类型，怎么实现？

解决方案：

因为 Python 不支持重载方法或函数，所以我们不能使用不同的签名定义 htmlize 的变体，也无法使用不同的方式处理不同的数据类型。在 Python 中，一种常见的做法是把 htmlize 变成一个分派函数，使用一串 if/elif/elif，调用专门的函数，如 htmlize_str、htmlize_int，等等。这样不便于模块的用户扩展，还显得笨拙：时间一长，分派函数 htmlize 会变得很大，而且它与各个专门函数之间的耦合也很紧密。

Python 3.4 新增的 functools.singledispatch 装饰器可以把整体方案拆分成多个模块，甚至可以为你无法修改的类提供专门的函数。使用 @singledispatch 装饰的普通函数会变成泛函数（generic function）：根据第一个参数的类型，以不同方式执行相同操作的一组函数。

上面的示例代码改造如下：

from functools import singledispatch
from collections import abc
import numbers
import html

@singledispatch
def htmlize(obj):
    content = html.escape(repr(obj))
    return '<pre>{}</pre>'.format(content)

@htmlize.register(str)
def _(text):
    content = html.escape(text).replace('\n', '<br>\n')
    return '<p>{0}</p>'.format(content)

@htmlize.register(numbers.Integral)
def _(n):
    return '<pre>{0} (0x{0:x})</pre>'.format(n)

@htmlize.register(tuple)
@htmlize.register(abc.MutableSequence)
def _(seq):
    inner = '</li>\n<li>'.join(htmlize(item) for item in seq)
    return '<ul>\n<li>' + inner + '</li>\n</ul>'

示例说明：

@singledispatch 标记处理 object 类型的基函数。各个专门函数使用 @«base_function».register(«type») 装饰。专门函数的名称无关紧要；下划线是个不错的选择，简单明了。为每个需要特殊处理的类型注册一个函数。numbers.Integral 是 int 的虚拟超类。可以叠放多个 register 装饰器，让同一个函数支持不同类型。只要可能，注册的专门函数应该处理抽象基类（如 numbers.Integral 和 abc.MutableSequence），不要处理具体实现（如 int 和 list）。这样，代码支持的兼容类型更广泛。例如，Python 扩展可以子类化 numbers.Integral，使用固定的位数实现 int 类型。

使用抽象基类检查类型，可以让代码支持这些抽象基类现有和未来的具体子类或虚拟子类。

singledispatch 机制的一个显著特征是，你可以在系统的任何地方和任何模块中注册专门函数。如果后来在新的模块中定义了新的类型，可以轻松地添加一个新的专门函数来处理那个类型。此外，你还可以为不是自己编写的或者不能修改的类添加自定义函数。

@singledispatch 不是为了把 Java 的那种方法重载带入 Python。在一个类中为同一个方法定义多个重载变体，比在一个函数中使用一长串 if/elif/elif/elif 块要更好。但是这两种方案都有缺陷，因为它们让代码单元（类或函数）承担的职责太多。@singledispath 的优点是支持模块化扩展：各个模块可以为它支持的各个类型注册一个专门函数。（很棒的特性，原来还能这样玩！）

另外，如何参数化装饰器，添加参数这层封装，只需再加一层函数的嵌套，示例:

def clock(morearg):
    def deco(func):
        def clocked(*args, **kwargs):
            '''I'm clocked function'''
            t = pc()
            result = func(*args, **kwargs)  # 执行主函数 重要
            elapsed = pc() - t # 主函数执行时间 重要

            # func的属性 重要
            name = func.__name__

            # 格式化参数  打印信息  不重要
            arg_str = ', '.join(repr(arg) for arg in args)
            if not kwargs:  
                print('[%0.8fs] %s(%s) -> %r' % (elapsed, name, arg_str, result))
            else:
                print('[%0.8fs] %s(%s ...) -> %r' % (elapsed, name, arg_str, result))

            return result
        
        return clocked
    return deco

第一层留给装饰器的参数，第二层是原来的被装饰函数的参数。

整理自《流畅的Python》第7章相关内容。

Python实现简单的装饰器

See Also

最近文章

标签

友情链接