[Typing-sig] Is it possible to annotate decorators that work on both functions and methods?

Feb. 26, 2023

      It appears to be impossible to perfectly annotation a decorator that works on both functions or methods.  And also, it's particularly hard to implement the __get__ method.  I propose adding typing.PositionalSubstitute to resolve this.

Consider:
from typing import Callable, Generic, Protocol, TypeVar, overload, Any

from typing_extensions import ParamSpec, Self, 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 __get__(self, instance: Any, owner: Any = None) -> Wrapped[P, V_co]:
        ...

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

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

def jit(f: Callable[..., Any]) -> Any:
    pass

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)  # Fails!
reveal_type(x.f)
reveal_type(g.__call__)

We can't seem to deal with the method case alongside the function case.  Here's the proposed solution:

class Wrapped(Protocol, Generic[P, V_co]):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> V_co:
        ...

    def __get__(self, instance: U, owner: Any = None
                ) -> Wrapped[PositionalSubstitute[P, U], V_co]:
        ...  # Much easier!

def jit(f: Callable[P, V_co]) -> Wrapped[P, V_co]:  # pyright: ignore
    pass  # No overloads!

The idea is that PositionalSubstitute takes a ParamSpec argument and returns it less that ParamSpec's first positional parameter.  It verifies that this removed parameter matches its second argument.

This could be generalized to substitute with tuples (or ParamSpec.args), or generalized to support multiple applications (taking a tuple argument as its second argument).

[Typing-sig] Is it possible to annotate decorators that work on both functions and methods?

Neil Girdhar