Python使用装饰器限定函数的参数范围

在特定的场景下,大量相似的函数需要对参数限定范围并报错,如果每个函数里面写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...

你可能感兴趣的:(Python)