python cookbook第三版学习笔记二十一:利用装饰器强制函数上的类型检查

时间:2024-01-01 15:11:33

在演示实际代码前,先说明我们的目标:能对函数参数类型进行断言,类似下面这样:

@typeassert(int, int)

... def add(x, y):

...     return x + y

add(2, 'hello')
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "contract.py", line 33, in wrapper
TypeError: Argument y must be <class 'int'>
 

我们可以自己实现这样一个装饰器。首先介绍下inspect.signature

sig=signature(test_func)

print(sig)

print(sig.parameters)运行结果:

(x, y, z=42)

OrderedDict([('x', <Parameter "x">), ('y', <Parameter "y">), ('z', <Parameter "z=42">)])signature方法将函数的参数签名信息转换成一个可调用对象

我们还可以使用sig.bind_partial()方法来执行从指定类型到名称的部分绑定。比如下面的绑定,这里将z这个参数绑定成了字节形式。返回结果是一个有序字典,这个字典会将参数名以函数签名中相同顺序映射到指定的类型值上面去

bound_types=sig.bind_partial(int,z=bytes)

print(bound_types.arguments)

OrderedDict([('x', <class 'int'>), ('z', <class 'bytes'>)])

另外一个是sig.bind方法。bind不允许忽略任何参数。Bind会将参数的值和变量名称绑定在一起,通过遍历就可以得出参数的名称和值。

bound_values=sig.bind(1,2,3)

print(bound_values)

for name,value in bound_values.arguments.items():

print(name,value)

运行结果:

<BoundArguments (x=1, y=2, z=3)>

x 1

y 2

z 3

通过bind_partial和bind方法我们就可以做到对参数的检查。bind_partial来对参数类型做强制绑定,然后用bind方法遍历出所有的参数进行参数核查。方法如下。

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

)

装饰器的代码如下:

def typeassert(*ty_args, **ty_kwargs):

def decorate(func):

# If in optimized mode, disable type checking

if not __debug__:

return func

# Map function argument names to supplied types

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)

# Enforce type assertions across supplied arguments

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,z=int)

def para_check(x,y,z=3):

print(x,y,z)

if __name__=="__main__":

para_check(1,2,'str')

运行结果:由于z这个参数赋的是一个字符,因此抛出异常。

Traceback (most recent call last):

File "D:/py_prj/test2/cookbook.py", line 266, in <module>

para_check(1,2,'str')

File "D:/py_prj/test2/cookbook.py", line 254, in wrapper

'Argument {} must be {}'.format(name, bound_types[name])

TypeError: Argument z must be <class 'int'>