如何“告訴”裝飾器 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: [...]