如何避免被装饰的Python函数丢失类型提示信息？

在VS Code中使用装饰器包装带类型提示的类方法时，被包装的方法会丢失类型提示，鼠标悬停无法查看原方法的参数和返回值类型。示例代码如下：

from typing import Callable
import functools

def decorate(function: Callable):
    @functools.wraps(function)
    def wrapper(object: "A", *args, **kwargs):
        return function(object, *args, **kwargs)
    return wrapper

class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

a = A()

print(a.g(2))

解决方案 #

• 用泛型类型标注装饰器：通过typing模块的ParamSpec和TypeVar，让类型检查器能正确推断原函数的参数和返回值类型，保留类型提示。
  修改后的代码：

from typing import Callable, ParamSpec, TypeVar
import functools

# 定义泛型参数，捕获原函数的参数签名和返回值类型
P = ParamSpec("P")
R = TypeVar("R")

def decorate(function: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(function)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        return function(*args, **kwargs)
    return wrapper

class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

a = A()
print(a.g(2))

这种方式让装饰器明确返回与原函数类型一致的包装函数，VS Code的类型检查器（如Pylance）可以正确识别原方法的类型提示。

• 针对实例方法的精准标注：如果装饰器仅用于类的实例方法，可以更明确地绑定实例类型，确保self参数的类型被正确保留：

from typing import Callable, TypeVar, ParamSpec
import functools

T = TypeVar("T")  # 绑定类实例类型
P = ParamSpec("P")
R = TypeVar("R")

def decorate(function: Callable[[T, *P.args], R]) -> Callable[[T, *P.args], R]:
    @functools.wraps(function)
    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
        return function(self, *args, **kwargs)
    return wrapper

class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

a = A()
print(a.g(2))

• 更新VS Code的Python插件：确保Pylance插件处于最新版本，新版本通常会修复更多类型推断的bug，提升装饰器类型提示的支持。

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