如何避免丢失装饰函数的类型提示

Question

我注意到，当包装一个带有类型提示的函数或方法时，使用Visual Studio 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))

当我在 Visual Studio Code 中悬停在名称 "g" 上时，我会失去类型提示信息。你知道有什么办法可以防止这种情况发生吗？

谢谢！

- rider45

Answer 1

Python 3.8（和3.9）最好的做法如下：

from __future__ import annotations
from functools import wraps
from typing import Any, Callable, TypeVar


T = TypeVar("T")


def decorate(function: Callable[..., T]) -> Callable[..., T]:
    @wraps(function)
    def wrapper(obj: A, *args: Any, **kwargs: Any) -> T:
        return function(obj, *args, **kwargs)
    return wrapper


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

这将保留返回类型信息，但不包括参数类型的详细信息。装饰器@wraps至少应该保持签名完整。如果装饰器应该是适用于A方法的通用装饰器，这已经是最好的了。

如果您希望更具体一些，您可以始终将函数类型限制为Callable[[A, B, C], T]，但那样装饰器就不再那么通用了。

---

如果您升级到Python 3.10 ，您将可以访问ParamSpec 。然后您可以执行以下操作：

from __future__ import annotations
from functools import wraps
from typing import Callable, ParamSpec, TypeVar


P = ParamSpec("P")
T = TypeVar("T")


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


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

这实际上保留了所有参数规范（正如名称所示）。

希望这有所帮助。

- Daniil Fajnberg

---

确实如此！非常感谢您提供的详细解释。 - rider45