functools.partial on function-decorator error

[YouJiacheng]

from functools import partial
from typing import TypeVar, Callable
from typing_extensions import ParamSpec, Protocol

T = TypeVar('T', covariant=True)
P = ParamSpec('P')

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

def jit(f: Callable[P, T]) -> Wrapped[P, T]:
    return f

def f(x: int, y: int):
    return x + y

partial_jit = partial(jit)
partial_jit(f)(0, 0)

error:

paramspec.py:18: error: Argument 1 to "partial" has incompatible type "Callable[[Callable[P, T]], Wrapped[P, T]]"; expected "Callable[..., Wrapped[P, T]]"
paramspec.py:19: error: Argument 1 to "__call__" of "Wrapped" has incompatible type "int"; expected "P.args"
paramspec.py:19: error: Argument 2 to "__call__" of "Wrapped" has incompatible type "int"; expected "P.args"
Found 3 errors in 1 file (checked 1 source file)

If change to:

from functools import partial
from typing import TypeVar, Callable
from typing_extensions import ParamSpec

T = TypeVar('T', covariant=True)
P = ParamSpec('P')

def jit(f: Callable[P, T]) -> Callable[P, T]:
    return f

def f(x: int, y: int):
    return x + y

partial_jit = partial(jit)
partial_jit(f)(0, 0)

error:

paramspec.py:15: error: Argument 1 has incompatible type "int"; expected "P.args"
paramspec.py:15: error: Argument 2 has incompatible type "int"; expected "P.args"
Found 2 errors in 1 file (checked 1 source file)

If change to:

from functools import partial
from typing import TypeVar, Callable
from typing_extensions import ParamSpec

T = TypeVar('T')
P = ParamSpec('P')


def jit(f: Callable[P, T]) -> Callable[P, T]:
    return f

def f(x: int, y: int):
    return x + y

partial_jit = partial(jit)
partial_jit(f)(0, 0)

error:

paramspec.py:16: error: Argument 1 has incompatible type "int"; expected "P.args"
paramspec.py:16: error: Argument 2 has incompatible type "int"; expected "P.args"
Found 2 errors in 1 file (checked 1 source file)

If change to:

from functools import partial
from typing import TypeVar

F = TypeVar('F')

def jit(f: F) -> F:
    return f

def f(x: int, y: int):
    return x + y

partial_jit = partial(jit)
partial_jit(f)(0, 0) # error here
jit(f)(0, 0) # no error here

error:

paramspec.py:13: error: "object" not callable
Found 1 error in 1 file (checked 1 source file)

[tmke8]

I think the main problem is the definition of partial in typeshed, which doesn't use ParamSpec at all. Unfortunately, the current type system can't really express all the behavior of partial. But if you use a definition like the following, it works:

from typing import Callable, Concatenate, Generic, ParamSpec, Protocol, TypeVar
from typing_extensions import reveal_type

P = ParamSpec("P")
R_co = TypeVar("R_co", covariant=True)
T = TypeVar("T")

class simple_partial(Generic[P, R_co, T]):
    """This version of `partial` can only set one argument and it has to be the first."""

    def __init__(self, __func: Callable[Concatenate[T, P], R_co], arg: T):
        self.func = __func
        self.arg = arg

    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R_co:
        return self.func(self.arg, *args, **kwargs)

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

def jit(flag: bool, f: Callable[P, T]) -> Wrapped[P, T]:
    return f

def f(x: int, y: int) -> int:
    return x + y

partial_jit = simple_partial(jit, True)
reveal_type(partial_jit)  # N: Revealed type is "simple_partial[[f: def (*P.args, **P.kwargs) -> T`-2], Wrapped[P`-1, T`-2], builtins.bool]"
f_wrapped = partial_jit(f)
reveal_type(f_wrapped)  # N: Revealed type is "Wrapped[[x: builtins.int, y: builtins.int], builtins.int]"
x = f_wrapped(0, 0)
reveal_type(x)  # N: Revealed type is "builtins.int"

[ilevkivskyi]

All four original examples work on current master, but some only with --new-type-inference.