装饰器+inspect的python入参检验

53 阅读 0 评论 35 点赞

我是靠谱客的博主顺利爆米花，最近开发中收集的这篇文章主要介绍装饰器+inspect的python入参检验，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

由于 python 不支持函数入参类型的限定，所以，对于 python 来说，入参合法性检测显得尤为重要。

def add(a: int, b: int) -> int:
    return a + b

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

-------------------------------------------------------------
Traceback (most recent call last):
  File "D:/MyPython/checker/checker_param.py", line 9, in <module>
    print(add(1, "2"))
  File "D:/MyPython/checker/checker_param.py", line 5, in add
    return a + b
TypeError: unsupported operand type(s) for +: 'int' and 'str'

python inspect

需要介绍一下 PEP 3107 之 Function Annotations。这个是 python3 才支持的一个特性，可以为函数的参数进行注解，如下所示。

上文的add函数的int也可以写成：

def add(x: 'integer', y: 'integer') -> 'the sum':
    return x + y

别人看到你的源码的时候，通过注解，就可以你这个参数应该传什么进去，实际上，这些注释是藏在 __annotations__ 字段里的，通过如下指令可以显示。

print(add.__annotations__)
{'x': 'integer', 'y': 'integer', 'return': 'the sum'}

字典 __annotations__ 里的值可以是任何 object，可以是 list、tuple、函数等。

直接上代码：

# coding=utf-8

import functools
import inspect

def check(name, value, checker):
    if isinstance(checker, (tuple, list, set)):
        return True in [check(name, value, sub_checker) for sub_checker in checker]
    elif checker is inspect._empty:
        return True
    elif checker is None:
        return value is None
    elif isinstance(checker, type):
        return isinstance(value, checker)
    elif callable(checker):
        result = checker(value)
        return result

def auto_type_checker(function):
    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        sig = inspect.signature(function)
        parameters = sig.parameters
        # fetch the argument name list.
        argument_list = list(parameters.keys())

        # fetch the parameter type list.
        type_list = [parameters[argument].annotation for argument in argument_list]

        # fetch the parameter value list.
        value_list = sig.bind(*args, **kwargs).arguments.values()

        # check the invalid argument, and raise the error.
        for argument, value, checker in zip(argument_list, value_list, type_list):
            if not check(argument, value, checker):
                raise Exception(f"arg '{argument}' required type:{checker},got type {type(value)}")

        # check the result.
        result = function(*args, **kwargs)
        checker = sig.return_annotation
        if not check('return', result, checker):
            raise Exception(['return'])

        # return the result.
        return result

    return wrapper

@auto_type_checker
def add(a: int, b: int) -> int:
    return a + b

if __name__ == '__main__':
    print(add.__annotations__)
    print(add(1, 4))

其中，auto_type_checker 是一个修饰器，在函数定义的时候调用即可。函数在声明的时候，如果需要进行入参合法性校验的话，就用如下语法为函数的输入输出指定 checker。

@auto_type_checker
def function(arg1: int, arg2, arg3: (int, float) = 0, arg4: lambda x: x > 0 = 1) -> list:
    return [arg1, arg2, arg3, arg4]

上述代码有 4 种 checker：

arg1 <class 'int'>：type 型 checker，如 arg1，auto_type_checker 会检测 arg1 是否是 int 型，如果不是，会抛出异常，而返回值必须是 list 型，否则也会抛出异常；
arg2 <class 'inspect._empty'>：不指定 checker，如 arg2，auto_type_checker 不会为 arg2 进行合法性校验；
arg3 (<class 'int'>, <class 'float'>)：tuple/list 型 checker，如 arg3，tuple 或 list 中的所有元素都会被当作 checker，当所有 checker 都无法通过校验，则抛出异常，上述代码中，arg3 允许整数或浮点数，0为默认值。
arg4 <function <lambda> at 0x0000020126AD8730>：函数型 checker，如 arg4，auto_type_checker 会将 arg4 带入到 checker，如果 checker 的返回值是 Fasle，则抛出异常，上述代码中，arg4 只接受大于 0 的数字，1为默认值。；

测试：

print(function(1, 2, 3, 4))
print(function(1, 2, 3.0, 4))
print(function(1, 2, 3, -4))
************************************
[1, 2, 3, 4]
[1, 2, 3.0, 4]
    print(function(1, 2, 3, -4))
  File "D:/MyPython/checker/checker_param.py", line 39, in wrapper
    raise Exception(f"arg '{argument}' required type:{checker},got type {type(value)}")
Exception: arg 'arg4' required type:<function <lambda> at 0x0000026EC42C8730>,got type <class 'int'>

更有的参考：亦可实现
typing —— 类型注解支持 — Python 3.11.2 文档

参考：