自定义参数类型断言装饰器

1|0代码

from inspect import signature
from functools import wraps


def typeassert(*ty_args, **ty_kwargs):
    def decorate(func):
        if not __debug__:
            return func
        sig = signature(func)  #获取函数签名
        bound_types = sig.bind_partial(*ty_args, **ty_kwargs).arguments

        @wraps(func)
        def wrapper(*args, **kwargs):
            bound_values = sig.bind(*args, **kwargs)
            for name, value in bound_values.arguments.items():
                if name in bound_types:
                    if not isinstance(value, bound_types[name]):
                        raise TypeError(
                            'Argument {} must be {}'.format(name, bound_types[name])
                        )
            return func(*args, **kwargs)
        return wrapper
    return decorate


@typeassert(int, int)
def add(x, y):
    return x+y


if __name__ == '__main__':
    print(add(1, 2))

2|0说明

这段代码定义了一个装饰器函数typeassert，其作用是对函数的参数类型进行断言检查。

首先，我们可以看到typeassert函数接收可变数量的位置参数ty_args和关键字参数ty_kwargs，这些参数用于指定被装饰函数的参数类型。

在decorate函数内部，首先进行了一个条件判断if not __debug__:。__debug__是Python内置的一个全局变量，
当使用命令行参数-O或-OO运行代码时，__debug__的值为False，否则为True。因此，这个判断语句的作用是当代码运行在优化模式下时，直接返回被装饰的函数，不进行类型断言检查。

如果不是在优化模式下运行，那么接下来就开始进行类型断言检查。首先通过signature函数获取被装饰函数的签名对象sig。签名对象用于描述函数的参数和返回值的类型信息。

然后调用sig.bind_partial(*ty_args, **ty_kwargs)方法，将传入的ty_args和ty_kwargs参数与被装饰函数的参数进行部分绑定。
这个部分绑定的目的是为了获取被装饰函数中指定了类型的参数。

接着，定义了一个内部函数wrapper，并使用@wraps(func)装饰器将其与被装饰的函数进行绑定。这个内部函数就是最终被返回的包装函数。

wrapper函数接收任意数量的位置参数args和关键字参数kwargs，这些参数是被装饰函数的实际参数。
然后，通过sig.bind(*args, **kwargs)方法将实际参数与被装饰函数的参数进行绑定。

接下来，使用一个循环遍历bound_values.arguments，其中bound_values.arguments是一个字典，包含了实际参数的绑定值。
在循环中，依次检查绑定值对应的参数是否在bound_types字典中。

如果参数名在bound_types字典中，说明该参数在ty_args或ty_kwargs中指定了类型。然后，通过isinstance函数检查绑定值的类型是否与指定的类型匹配。
如果类型不匹配，则抛出TypeError异常，提示参数类型错误。

如果所有参数类型都通过了断言检查，那么最后返回调用被装饰函数的结果。