使用 wrapt 模块编写更通用的装饰器

使用 wrapt 模块编写更通用的装饰器

# wrapt是第三方包，使用需要安装
pip install wrapt

1. 常见写法：

import random


def provide_number(min_num, max_num):
    """装饰器：随机生成一个在 [min_num, max_num] 范围的整数，追加为函数的第一个位置参数
    """
    def wrapper(func):
        def decorated(*args, **kwargs):
            num = random.randint(min_num, max_num)
            # 将 num 作为第一个参数追加后调用函数
            return func(num, *args, **kwargs)
        return decorated
    return wrapper
    


@provide_number(1, 100)
def print_random_number(num):
    print(num)

# 输出 1-100 的随机整数
# OUTPUT: 72
print_random_number()

@provide_number 装饰器功能看上去很不错，但它有着我在前面提到的两个问题：**嵌套层级深、无法在类方法上使用。**如果直接用它去装饰类方法，会出现下面的情况：

class Foo:
    @provide_number(1, 100)
    def print_random_number(self, num):
        print(num)

# OUTPUT: <__main__.Foo object at 0x104047278>
Foo().print_random_number()

Foo 类实例中的 print_random_number 方法将会输出类实例 self ，而不是我们期望的随机数 num。

2. 通用写法：

import wrapt

def provide_number(min_num, max_num):
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        # 参数含义：
        #
        # - wrapped：被装饰的函数或类方法
        # - instance：
        #   - 如果被装饰者为普通类方法，该值为类实例
        #   - 如果被装饰者为 classmethod 类方法，该值为类
        #   - 如果被装饰者为类/函数/静态方法，该值为 None
        #
        # - args：调用时的位置参数（注意没有 * 符号）
        # - kwargs：调用时的关键字参数（注意没有 ** 符号）
        #
        num = random.randint(min_num, max_num)
        # 无需关注 wrapped 是类方法或普通函数，直接在头部追加参数
        args = (num,) + args
        return wrapped(*args, **kwargs)
    return wrapper
    
<... 应用装饰器部分代码省略 ...>
    
# OUTPUT: 48
Foo().print_random_number()

使用 wrapt 模块编写的装饰器，相比原来拥有下面这些优势：

• 嵌套层级少：使用 @wrapt.decorator 可以将两层嵌套减少为一层
• 更简单：处理位置与关键字参数时，可以忽略类实例等特殊情况
• 更灵活：针对 instance 值进行条件判断后，更容易让装饰器变得通用

参考文档：https://github.com/piglei/one-python-craftsman/blob/master/zh_CN/8-tips-on-decorators.md