How can we annotate a decorator that works with functions and methods?

[NeilGirdhar]

It appears to be impossible to perfectly annotate a decorator that works on both functions or methods. In particular, it's practically impossible to implement the __get__ method.

It is also currently impossible to natively (without plugins) annotate functools.partial. Even a MyPy plugin has proved to be very difficult.

I propose adding typing.PartialApplication to:

• easily implement __get__,
• get us halfway to natively implementing functools.partial.

Consider trying to create a decorator jit that works with both bare functions and methods. The problem is that in the method case, it has to respond to __get__ and strip off the first argument. It seems that we can only do this with Concatenate:

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 WrappedMethod(Protocol, Generic[S, P, V_co]):
    def __call__(self: S, *args: P.args, **kwargs: P.kwargs) -> V_co:
        ...

    @overload
    def __get__(self, instance: None, owner: Any = None) -> Self:
        ...

    @overload
    def __get__(self, instance: S, owner: Any = None) -> Wrapped[P, V_co]:
        ...

# this overload can only be hit if there is a positional parameter.  It responds to `__get__` by throwing that parameter out.  
@overload
def jit(f: Callable[Concatenate[U, P], V_co]) -> MethodWrapped[U, P, V_co]:  # type: ignore
    ...

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

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

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[PartialApplication[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 PartialApplication 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

• support multiple applications (taking a tuple or ParamSpec.args as its second argument), or
• support applying keyword arguments (taking a TypedDict or ParamSpec.kwargs as its third argument).

Idealized example with partial (might need some tweaks)

P = ParamSpec('P')
Q = ParamSpec('P')
R = TypeVar('R', covariant=True)
class partial(Generic[P, Q, R]):
  S: TypeAlias = PartialApplication(P, Q.args, Q.kwargs)
  def __init__(self, f: Callable[P, R], /, *args: Q.args, **kwargs: Q.kwargs): ...
  def __call__(self, /, *args: S.args, **kwargs: S.kwargs) -> R: ...

[tmke8]

Isn't the problem with decorators for methods more that the type checker doesn't know it is supposed to bind the self argument?

If you have a decorator like this:

def decorator(f: Callable[P, R]) -> Callable[P, R]:
    f.my_attribute = 4  # type error
    return f

class C:
    @decorator
    def m(self, x: int) -> int:
        return 2 * x

c = C()
reveal_type(C.m)
reveal_type(c.m)
print(c.m(4))
print(c.m.my_attribute)  # type error

then type checkers can infer the signature of m() correctly, without using descriptors. Though in this case, we can't express that the decorated method has additional attributes.

It seems to me what's needed is a way to define subclasses of functions, so that they still get bound as methods but can have additional attributes.

[tmke8]

The obvious attempt to solve the above is this:

class FuncWithAttribute(Protocol[P, R]):
    my_attribute: int
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: ...

def decorator(f: Callable[P, R]) -> FuncWithAttribute[P, R]: ...

class C:
    @decorator
    def m(self, x: int) -> int:
        return 2 * x

c = C()
print(c.m(4))  # type error
print(c.m.my_attribute)  # OK now

but then the function isn't bound anymore according to type checkers.

[NeilGirdhar]

but then the function isn't bound anymore according to type checkers.

Right, you need to implement the descriptor protocol, which means adding the __get__ method.

[tmke8]

Right, but what I wanted to say is that the implementation doesn't use the descriptor protocol (see for example in jax: https://github.com/google/jax/blob/7b6321cc09a975417fc6e7dbde4e4cc168009acd/jax/_src/api.py#L504-L546 – they just return a function to which they assign attributes; they also define a protocol here: https://github.com/google/jax/blob/7b6321cc09a975417fc6e7dbde4e4cc168009acd/jax/_src/stages.py#L654-L677 but it's only for type checking), so why should the type annotations have to use the descriptor protocol? It seems to me, a more elegant solution would be to teach type checkers which objects will get bound as methods.

[NeilGirdhar]

see for example in jax:

Their implementation doesn't preserve the parameter list.

why should the type annotations have to use the descriptor protocol?

Because if they don't then you can't apply the decorator to a method while preserving the parameter list.

a more elegant solution would be to teach type checkers which objects will get bound as methods.

I think that kind of "hack" would be far less elegant and flexible. What if, for example, the decorator wants to add a parameter? Or what if it wants to add a copy of the self parameter? Or what if you inherit from the decorator and want to add behavior to __get__ in a subclass.

And anyway, why shouldn't you be able to annotate __get__? It just seems like an unnecessary hole in the type system.