Add PartialApplication

[NeilGirdhar]

Pitch

Add typing.PartialApplication to facilitate the implementation of:

• __get__, and
• functools.partial,

both of which are practically impossible to natively (without plugins) annotate.

The __get__ method is currently handled internally by type checkers, and a MyPy plugin for partial has proved to be very difficult.

Proposal

I created a discussion, but I wanted to flesh this out as a new feature:

The idea is that PartialApplication takes three parameters:

• a ParamSpec parameter P,
• a tuple parameter T, and
• a dictionary parameter D (defaulting to an empty dictionary).

It returns a new ParamSpec with all the arguments of P after removing

• the first len(T) positional parameters,
• the named keyword parameters from D.

It verifies that this removed parameters are all supertypes of the corresponding arguments, or else returns a type error.

Partial case study

An 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: ...

Thus, calling partial(f, ...) would check the parameters, and produce a __call__ method with the right signature.

JIT example

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]) -> WrappedMethod[U, P, V_co]:
    ...

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

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
                ) -> Callable[PartialApplication[P, tuple[U]], V_co]:
        ...  # Much easier!

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

[erictraut]

Pyright already has very complete support for functools.partial. Mypy also has support for this, although (as you've pointed out), it is incomplete in some areas. So I don't think that functools.partial is a good justification for adding a complex new type feature to the type system.

If we were to replace pyright's custom logic for functools.partial with your proposal, it would be a big step backward in terms of functionality (e.g. loss of bidirectional type inference when evaluating arguments) and specificity of error messages. So I don't see any upsides to adopting this proposal for functools.partial in pyright.

I agree that there's a general problem with ParamSpec not dealing well with the difference between methods and functions, but I'm not convinced that PartialApplication is the right solution to that problem. Then again, I haven't thought deeply about what a better solution might entail.

[NeilGirdhar]

Pyright already has very complete support for functools.partial. Mypy also has support for this, although (as you've pointed out), it is incomplete in some areas

Right, it's incomplete. MyPy doesn't keep track of the parameters after application.

loss of bidirectional type inference when evaluating arguments

Why can't this be done with PartialApplication?

and specificity of error messages

I guess I don't see why PartialApplication wouldn't do exactly what you're doing for partial—including return the exact same error messages? It seems like this proposal would just ask Pyright to expose to the typing user whatever you're doing with partial.

I don't think that functools.partial is a good justification for adding a complex new type feature to the type system.

Given that Pyright already supports partial, I understand your point. How can I annotate the __get__ method for Pyright though?

Also, there are custom versions of partial. For example, jax.tree_util.Partial, which would benefit from accessing the same logic that's baked into Pyright's implementation of partial.

[NeilGirdhar]

In Python, there are three major flavours of callables in Python. Decorators can return any of these, and there's no way to specify them currently. Instead, type checkers essentially guess that the decorator returns the same flavour that was passed in. Were this proposal accepted, we could annotate these explicitly, and declare what the decorator is doing:

from collections.abc import Callable
from typing import Generic


from __future__ import annotations
from collections.abc import Callable
from typing import Any, Generic, ParamSpec, Protocol, TypeVar, overload

P = ParamSpec('P')
U = ParamSpec('U')
R_co = TypeVar('R_co', covariant=True)


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

    @overload
    def __get__(self, instance: None, owner: Any = None) -> OrdinaryCallable[P, R_co]:
        ...
    @overload
    def __get__(self, instance: U, owner: Any = None
                ) -> Callable[PartialApplication[P, tuple[U]], R_co]:  # Bind instance!
        ...
    def __get__(self, instance: Any, owner: Any = None) -> Callable[..., R_co]:
        ...


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

    @overload
    def __get__(self, instance: None, owner: Any = None) -> StaticMethod[P, R_co]:
        ...
    @overload
    def __get__(self, instance: Any, owner: Any = None
                ) -> OrdinaryCallable[P, R_co]:  # Never bind!
        ...
    def __get__(self, instance: Any, owner: Any = None) -> Callable[..., R_co]:
        ...


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

    @overload
    def __get__(self, instance: None, owner: U = None
                ) -> OrdinaryCallable[PartialApplication[P, tuple[U]], R_co]:  # Bind owner!
        ...
    @overload
    def __get__(self, instance: U, owner: type[U] = None
                ) -> OrdinaryCallable[PartialApplication[P, tuple[type[U]]], R_co]:  # Bind instance!
        ...
    def __get__(self, instance: Any, owner: Any = None) -> Callable[..., R_co]:
        ...

Then, we could do something like

def my_decorator(func: OrdinaryCallable[P, R]) -> OrdinaryCallable[P, R]:

Or StaticMethod or ClassMethod, or some overloaded combination. This would essentially bring to the surface the arcane magic that various type checkers are currently doing.

[Gobot1234]

I've been thinking about this a lot recently and I think it is the best approach to solving this problem. PEP 612's lack of a way to solve this problem brings me great displeasure (see #946 for another concrete example of where this'd be useful), but I understand that there's no feasible way to solve this with just Concatenate so I think the addition of a special form this is the best way to allow for this to be expressible.

I see this problem come up very regularly with people complaining about functools.cache/their own implementation of a similar concept not persevering parameters and also with partial not being fully expressible in the type system. This would also help with an idea for making FunctionType and MethodType types subscriptable to allow for more specific types that just Callable for some things (see #1410).