如何“告诉”装饰器 function 返回类型与 Python 中被包裹的 function 相同？

Question

我有一个装饰器，其返回类型与包装的 function 返回类型相同。 这是一个简单的例子。 该代码按原样运行，但是，它不为 output 值res提供类型提示。 有没有办法指定my_decorator返回与包装的 function foo相同的类型（例如，两者都返回Foo类型的变量）

@dataclasses.dataclass
class Foo:
    x: str

def my_decorator(func):
    def wrapped(x):
        print("decorated")
        return func(x)
    return wrapped


@my_decorator
def foo(x: int) -> Foo:
    return Foo(str(x))


if __name__ == '__main__':
    res = foo(1) # no type hinting for res

Answer 1

您将使用ParamSpec和TypeVar ，这样您就可以参考参数并返回您传递的 function 的类型。（文档： ParamSpec ， TypeVar ）

所以Callable[P, R]是你的装饰器接受的参数，也是它的返回类型。 如果它们在运行时不同，那么您的 static 类型检查器将捕获它。 内部 function 也应指定它采用P并输出R 。 由于ParamSpec将捕获所有位置和关键字 arguments，您需要使用 arguments 声明内部 function： *args: P.args, **kwargs: P.kwargs ，因此您的代码最终将如下所示：

from typing import Callable, ParamSpec, TypeVar

@dataclasses.dataclass
class Foo:
    x: str

P = ParamSpec('P')
R = TypeVar('R')

def my_decorator(func: Callable[P, R]) -> Callable[P, R]:
    def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
        print("decorated")
        return func(*args, **kwargs)
    return wrapped

@my_decorator
def foo(x: int) -> Foo:
    return Foo(str(x))

使用mypy测试：

foo(5) # Success: [...]