如何让Python装饰器返回与被装饰函数相同的类型？

要让装饰器保留被装饰函数的参数和返回类型信息，需要给装饰器添加泛型类型注解，让类型检查器能正确推断返回值类型。以下是两种可行方案：

方案1：兼容Python 3.9及以下版本 #

使用TypeVar和Callable定义泛型装饰器，适用于参数固定的场景：

from dataclasses import dataclass
from typing import TypeVar, Callable

T = TypeVar('T')  # 定义泛型类型变量，代表返回类型

@dataclass
class Foo:
    x: str

def my_decorator(func: Callable[[int], T]) -> Callable[[int], T]:
    def wrapped(x: int) -> T:
        print("decorated")
        return func(x)
    return wrapped

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

if __name__ == '__main__':
    res = foo(1)  # 现在会正确提示res为Foo类型

方案2：Python 3.10+ 通用方案（推荐） #

使用ParamSpec捕获任意参数组合，让装饰器支持任意签名的函数：

from dataclasses import dataclass
from typing import TypeVar, ParamSpec, Callable

P = ParamSpec('P')  # 捕获函数的参数类型
T = TypeVar('T')    # 捕获函数的返回类型

@dataclass
class Foo:
    x: str

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

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

if __name__ == '__main__':
    res = foo(1)  # 正确提示res为Foo类型

原理说明 #

原来的装饰器没有类型注解，类型检查器无法关联wrapped函数和原函数的类型信息。通过泛型TypeVar和ParamSpec，我们让装饰器的输入和输出类型严格绑定到被装饰函数的参数与返回类型，从而保留完整的类型提示。

内容的提问来源于stack exchange，提问作者motam79