...

Python3使用装饰器实现参数类型检查

import typing
from typing import Union
from functools import wraps


def merge_args(varnames: tuple, args: tuple, kwargs: dict) -> dict:
    merged_kwargs = kwargs.copy()
    if args:
        merged_kwargs.update(dict(zip(varnames, args)))
    return merged_kwargs


def check_arg_type(arg_name, arg_value, arg_type):
    origin_arg_type = arg_type
    if isinstance(arg_type, typing._GenericAlias):
        arg_type = typing.get_args(arg_type)
    if not isinstance(arg_value, arg_type):
        raise TypeError(f'参数 "{arg_name}" 不符合注释类型 {origin_arg_type}')


def type_check(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        types_hint = func.__annotations__  # dict or None
        if types_hint:
            merged_kwargs = merge_args(func.__code__.co_varnames, args, kwargs)
            for arg_name, arg_value in merged_kwargs.items():
                if arg_name in types_hint:
                    arg_type = types_hint[arg_name]
                    check_arg_type(arg_name, arg_value, arg_type)

        return func(*args, **kwargs)

    return wrapper


@type_check
def add(a: Union[int, float], b: Union[int, float]):
    s = a + b
    return s


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

运行结果如下:

Traceback (most recent call last):
  File "/Users/superhin/Projects/试验/d01/18_装饰器_参数检查.py", line 44, in <module>
    print(add(1, '2'))
  File "/Users/superhin/Projects/试验/d01/18_装饰器_参数检查.py", line 30, in wrapper
    check_arg_type(arg_name, arg_value, arg_type)
  File "/Users/superhin/Projects/试验/d01/18_装饰器_参数检查.py", line 18, in check_arg_type
    raise TypeError(f'参数 "{arg_name}" 不符合注释类型 {origin_arg_type}')
TypeError: 参数 "b" 不符合注释类型 typing.Union[int, float]
posted @ 2023-06-19 16:09  韩志超  阅读(79)  评论(0编辑  收藏  举报