How to type-annotate a decorator that has an optional parameter?

Question

Here's the exact function that I'm trying to type-annotate correctly:

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)

While I'm not getting any error when putting it through mypy, I'm not convinced that I did the right thing, nor am I sure how I could go about proving it.

Answer 1

Your code typechecks but probably it does not do what you want, because you are returning an Union.

To check what type mypy infers for some variable you can use 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]'

Assuming we want hello1 to be a function that returns an int (i.e. def () -> builtins.int) we need to try something else.

Simple strategy

The simplest thing is to always ask the user of throttle to "call the decorator" even if she/he is not overriding any arguments:

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'

This already works and is very simple.

Use typing.overload

In case the previous solution is not acceptable, you can use 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'

You can learn more on how to use overload by reading its official documentation, and mypy's documentation on Function overloading.