util.py 文件源码

def autocast_decorator(type_hint, fix_arg_func):
    """
    Decorator which will invoke fix_arg_func for any
    arguments annotated with type_hint. The decorated
    function will then be called with the result.

    :param type_hint: A PEP484 type hint
    :param fix_arg_func: Function to invoke

    :return: decorator
    """
    @decorator
    def wrapper(wrapped, instance, args, kwargs):
        hinted_args = names = None
        cache_key = '%s-%s-%s' % (wrapped.__class__.__name__,
                                  wrapped.__name__, str(type_hint))

        if cache_key in AUTOCAST_CACHE:
            hinted_args, names = AUTOCAST_CACHE[cache_key]
        else:
            sig = inspect.signature(wrapped)
            names = list(sig.parameters.keys())
            hinted_args = [x[0] for x in typing.get_type_hints(wrapped).items() \
                    if x[1] == type_hint or x[1] == typing.Union[type_hint, None]]
            AUTOCAST_CACHE[cache_key] = hinted_args, names

        if len(hinted_args) == 0:
            raise ValueError("No arguments with %s hint found" % type_hint)

        new_args = list(args)
        for hinted_arg in hinted_args:
            if hinted_arg in kwargs:
                kwargs[hinted_arg] = fix_arg_func(kwargs[hinted_arg])

            elif hinted_arg in names:
                idx = names.index(hinted_arg)
                if idx < len(new_args):
                    new_args[idx] = fix_arg_func(new_args[idx])

        return wrapped(*new_args, **kwargs)

    return wrapper