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

代码

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))

说明

这段代码定义了一个装饰器函数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异常,提示参数类型错误。

如果所有参数类型都通过了断言检查,那么最后返回调用被装饰函数的结果。
posted @ 2023-07-11 08:57  我在路上回头看  阅读(17)  评论(0编辑  收藏  举报