Python中如何为支持部分应用的函数/方法做类型注解？

你需要实现一个装饰器jit，它既能处理普通函数，也能处理类方法：当装饰器作用于方法时，通过__get__绑定实例后，会自动把实例作为第一个参数部分应用；当处理普通函数时，要保留原函数的完整签名（包括关键字参数调用支持）。下面是修正后的类型注解实现：

修正后的完整代码 #

from typing import Callable, Generic, Protocol, TypeVar, overload, Any
from typing_extensions import ParamSpec, Concatenate

V_co = TypeVar("V_co", covariant=True)
U = TypeVar("U", contravariant=True)
P = ParamSpec("P")

class Wrapped(Protocol, Generic[P, V_co]):
    # 移除仅限位置参数标记`/`，允许关键字参数调用
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> V_co:
        ...

class MethodWrapped(Protocol, Generic[U, P, V_co]):
    # 方法未绑定时，需要接收实例作为第一个参数
    def __call__(self, u: U, *args: P.args, **kwargs: P.kwargs) -> V_co:
        ...

    # 实现描述符协议，绑定实例后返回已绑定的可调用对象；未绑定时返回自身
    def __get__(self, instance: U | None, owner: type[U] | None = None) -> Wrapped[P, V_co] | Self:
        ...

@overload
def jit(f: Callable[Concatenate[U, P], V_co]) -> MethodWrapped[U, P, V_co]:
    pass

@overload
def jit(f: Callable[P, V_co]) -> Wrapped[P, V_co]:
    pass

def jit(f: Callable[..., Any]) -> Any:
    # 实际装饰器逻辑（此处仅做类型注解，具体实现略）
    return f


class X:
    @jit
    def f(self, x: int) -> None:
        pass

@jit
def g(x: int, y: float) -> None:
    pass

x = X()
x.f(3)          # 类型检查通过
x.f(x=3)        # 类型检查通过
g(3, 4.2)       # 类型检查通过
g(x=3, y=4.2)   # 现在类型检查通过！

reveal_type(x.f)       # 应为 Wrapped[[int], None]
reveal_type(g.__call__) # 应为 Wrapped[[int, float], None]

关键修正点说明 #

1. 移除Wrapped协议的仅限位置参数标记
   原代码中__call__的/会强制参数只能按位置传递，导致普通函数g无法使用关键字参数调用。移除后，类型检查器会允许符合原函数签名的位置和关键字参数调用。

2. 完善__get__的返回类型
   原__get__仅返回Wrapped[P, V_co]，但当instance为None时（比如直接访问X.f），应该返回MethodWrapped自身。修正后的返回类型Wrapped[P, V_co] | Self覆盖了未绑定和已绑定两种场景，类型推导更准确。

3. 消除不必要的忽略注释
   原代码中的# type: ignore和# pyright: ignore已无需保留，修正后的类型注解能被Pyright、Mypy等类型检查器正确识别。

核心原理 #

• ParamSpec：捕获函数的完整参数签名（包括位置和关键字参数），保证装饰后的函数保留原参数的类型和传递方式。
• Concatenate：拼接实例类型和方法的参数签名，实现方法未绑定时的类型推导（比如X.f的签名是Callable[[X, int], None]，绑定实例后变为Callable[[int], None]）。
• 描述符协议（__get__）：让装饰后的方法能正确处理实例绑定，类型检查器会根据__get__的返回类型推导绑定后的函数签名。

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