python - 如何对具有可选参数的装饰器进行类型注释？

Question

这是我尝试正确输入注释的确切函数:

F = TypeVar('F', bound=Callable[..., Any])

def throttle(_func: Optional[F] = None, *, rate: float = 1) -> Union[F, Callable[[F], F]]:
    """Throttles a function call, so that at minimum it can be called every `rate` seconds.

    Usage::

        # this will enforce the default minimum time of 1 second between function calls
        @throttle
        def ...

    or::

        # this will enforce a custom minimum time of 2.5 seconds between function calls
        @throttle(rate=2.5)
        def ...

    This will raise an error, because `rate=` needs to be specified::

        @throttle(5)
        def ...
    """

    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            time.sleep(rate)
            return func(*args, **kwargs)

        return cast(F, wrapper)

    if _func is None:
        return decorator
    else:
        return decorator(_func)

虽然我在通过 mypy 传递它时没有收到任何错误，但我不相信我做的是正确的，我也不确定我该如何证明它。

Answer 1

您的代码进行了类型检查，但它可能没有按照您的要求执行，因为您返回的是 Union。

要检查 mypy 为某些变量推断的类型，您可以使用 reveal_type。

# Note: I am assuming you meant "throttle" and so I changed your spelling
def throttle1(
    _func: Optional[F] = None, *, rate: float = 1.0
) -> Union[F, Callable[[F], F]]:
    # code omitted


@throttle1
def hello1() -> int:
    return 42


reveal_type(hello1) # Revealed type is 'Union[def () -> builtins.int, def (def () -> builtins.int) -> def () -> builtins.int]'

假设我们希望 hello1 成为一个返回 int 的函数(即 def () -> builtins.int)，我们需要尝试其他方法。

简单策略

最简单的事情是始终要求 throttle 的用户“调​​用装饰器”，即使她/他没有覆盖任何参数:

def throttle2(*, rate: float = 1.0) -> Callable[[F], F]:
    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            time.sleep(rate)
            return func(*args, **kwargs)

        return cast(F, wrapper)

    return decorator


@throttle2() # Note that I am calling throttle2 without arguments
def hello2() -> int:
    return 42

reveal_type(hello2) # Revealed type is 'def () -> builtins.int'


@throttle2(rate=2.0)
def hello3() -> int:
    return 42

reveal_type(hello3) # Revealed type is 'def () -> builtins.int'

这已经可以工作并且非常简单。

使用typing.overload

如果前面的解决方案 Not Acceptable ，您可以使用overload。

# Matches when we are overriding some arguments
@overload
def throttle3(_func: None = None, *, rate: float = 1.0) -> Callable[[F], F]:
    ...

# Matches when we are not overriding any argument
@overload
def throttle3(_func: F) -> F:
    ...


def throttle3(
    _func: Optional[F] = None, *, rate: float = 1.0
) -> Union[F, Callable[[F], F]]:
    # your original code goes here


@throttle3 # Note: we do not need to call the decorator
def hello4() -> int:
    return 42


reveal_type(hello4) # Revealed type is 'def () -> builtins.int'


@throttle3(rate=2.0)
def hello5() -> int:
    return 42


reveal_type(hello5) # Revealed type is 'def () -> builtins.int'

您可以通过阅读 its official documentation 了解有关如何使用 overload 的更多信息, 和 mypy's documentation on Function overloading .