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

Question

我注意到，当包装具有某种类型提示的函数或方法时，当我使用 Visual Studio 代码进行编码时，包装的方法会丢失类型提示信息。

例如使用此代码：

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 代码中悬停在名称 g 上时，我丢失了类型提示信息。 你知道防止这种情况发生的方法吗？

真挚地

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 的方法具有通用性，那么这在 IMO 中就已经很好了。

如果你想让它更具体，你总是可以将函数类型限制为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}"

这实际上保留了（顾名思义）所有参数规范。

希望这可以帮助。