可以使用或不使用参数的类型装饰器

Question

我有一个可以不带参数或带参数（所有字符串）调用的装饰器：

@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

尽管阅读了 related section of the doc of mypy.

，但我很难提供适当的类型提示，以便类型检查器能够正确验证装饰器的使用

这是我目前尝试过的方法：

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 

    # @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

这导致 mypy 出现以下错误：

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

我也有一个错误 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

我正在使用 python 3.10.4，mypy 0.960，pyright 1.1.249。

Answer 1

问题来自第一个过载（我应该读两遍 pyright 消息！）：

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

此重载接受名为 arg 的关键字参数，而实现不接受！

当然，对于使用 @decorator 符号的装饰器来说，这并不重要，但如果这样调用就可以了：fct2 = decorator(arg=fct).

Python >= 3.8

解决该问题的最佳方法是更改第一个重载，使 arg 成为 positional-only parameter（因此不能用作关键字参数）：

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

支持 Python < 3.8

由于 positional-only 参数随 Python 3.8 一起提供，我们无法根据需要更改第一个重载。

相反，让我们更改实现以允许 **kwargs 参数（另一种可能性是添加关键字 arg 参数）。但是现在我们需要在代码实现中妥善处理，例如：

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

    # rest of the implementation here

可以使用或不使用参数的类型装饰器

Typing decorators that can be used with or without arguments

python

python-decorators

mypy

python-typing

pyright

Python >= 3.8

支持 Python < 3.8