在特定的场景下,大量相似的函数需要对参数限定范围并报错,如果每个函数里面写assertion可能会看起来比较冗杂,所以考虑用装饰器,看起来更规整点。
主要方案是利用函数的method: func.__code__.co_varnames
提取参数/默认值等实现。
# 规则
def not_none(x, i):
assert x is not None, f'{i} should not take None...'
return
# 装饰器,不支持叠加多个装饰器,外层的装饰器在func.__code__.co_varnames读取的参数将是底层装饰器的参数
def para_checker(para_single_rule, paras):
def wrapper(func):
def __(*args, **kwargs):
for i in paras:
func_para = func.__code__.co_varnames[:func.__code__.co_argcount]
if i in func_para:
idx = func_para.index(i)
# kwargs check
if i in kwargs.keys():
para_single_rule( kwargs[i], i)
# positional check
elif args.__len__() - 1 >= idx:
para_single_rule(args[idx], i)
# default para check
else:
pos = idx - func_para.__len__()
if (func.__defaults__.__len__() >= pos):
para_single_rule(func.__defaults__[pos], i)
return func(*args, **kwargs)
return __
return wrapper
# 以下仅对a,b作不能为None的规则限制
@para_checker(not_none, ['a','b'])
def f(a, e=1, c=None):
b = a+d
g = e+b
return g
f(None)
>>> AssertionError: a should not take None...