CODEWITHSUNDEEP
pythonpython-decoratorsmypypython-typingpyright
rogdham
rogdham

Reputation: 228

Typing decorators that can be used with or without arguments

I have a decorator that can be called either without or with arguments (all strings):

@decorator
def fct0(a: int, b: int) -> int:
    return a * b


@decorator("foo", "bar")  # any number of arguments
def fct1(a: int, b: int) -> int:
    return a * b

I am having a hard time providing appropriate type hints so that type checkers will be able to properly validate the usage of the decorator, despite having read the related section of the doc of mypy.

Here is what I have tried so far:

from typing import overload, TypeVar, Any, Callable

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

@overload
def decorator(arg: F) -> F:
    ...

@overload
def decorator(*args: str) -> Callable[[F], F]:
    ...

def decorator(*args: Any) -> Any:
    # python code adapted from https://stackoverflow.com/q/653368

    # @decorator -> shorthand for @decorator()
    if len(args) == 1 and callable(args[0]):
        return decorator()(args[0])

    # @decorator(...) -> real implementation
    def wrapper(fct: F) -> F:
        # real code using `args` and `fct` here redacted for clarity
        return fct

    return wrapper

Which results in the following error from mypy:

error: Overloaded function implementation does not accept all possible arguments of signature 1

I also have an error with pyright:

error: Overloaded implementation is not consistent with signature of overload 1
  Type "(*args: Any) -> Any" cannot be assigned to type "(arg: F@decorator) -> F@decorator"
    Keyword parameter "arg" is missing in source

I am using python 3.10.4, mypy 0.960, pyright 1.1.249.

Upvotes: 5

Views: 1124

Answers (1)

rogdham
rogdham

Reputation: 228

The issue comes from the first overload (I should have read the pyright message twice!):

@overload
def decorator(arg: F) -> F:
    ...

This overload accepts a keyword parameter named arg, while the implementation does not!

Of course this does not matter in the case of a decorator used with the @decorator notation, but could if it is called like so: fct2 = decorator(arg=fct).

Python >= 3.8

The best way to solve the issue would be to change the first overload so that arg is a positional-only parameter (so cannot be used as a keyword argument):

@overload
def decorator(arg: F, /) -> F:
    ...

With support for Python < 3.8

Since positional-only parameters come with Python 3.8, we cannot change the first overload as desired.

Instead, let's change the implementation to allow for a **kwargs parameter (an other possibility would be to add a keyword arg parameter). But now we need to handle it properly in the code implementation, for example:

def decorator(*args: Any, **kwargs: Any) -> Any:
    if kwargs:
        raise TypeError("Unexpected keyword argument")

    # rest of the implementation here

Upvotes: 4

Related Questions