Search code examples
python-typingmypypython-3.10

Correct annotation for "apply" function


In Python-3.10 (it must be this version) I want to add better annotation for my apply function:

from typing import TypeVar, Callable, Sequence, Any

T = TypeVar('T')

def apply(fn: Callable[..., T], vals: Sequence[Any]) -> T:
    return fn(*vals)

def f(a: int, b: int) -> int:
    return a + b

print(apply(f, (1, 2)))

Because in this code, the type annotation "hides" the connection between vals and fn arguments, so mypy can't check types and number of arguments correctly.

What have I tried:

1. using ParamSpec

from typing import ParamSpec

P = ParamSpec('P')

def apply(fn: Callable[P, T], vals: P) -> T:
    return fn(*vals)

which gives this error:

tst_apply.py:6: error: Invalid location for ParamSpec "P"  [valid-type]
tst_apply.py:6: note: You can use ParamSpec as the first argument to Callable, e.g., "Callable[P, int]"
tst_apply.py:7: error: Too few arguments  [call-arg]
Found 2 errors in 1 file (checked 1 source file)

2. using ParamSpec.args

from typing import ParamSpec

P = ParamSpec('P')

def apply(fn: Callable[P, T], vals: P.args) -> T:
    return fn(*vals)

which gives this error:

tst_apply.py:7: error: Too few arguments  [call-arg]
Found 1 error in 1 file (checked 1 source file)

3. using TypeVarTuple as suggested in Daraan's answer Revision 1:

Even though TypeVarTuple is available only from python-3.11, I can import it from typing_extensions, but it still doesn't work - it works when I change vals: *Ts to *vals: *Ts but that's unusable for me:

Ts = TypeVarTuple('Ts')

def apply(fn: Callable[[*Ts], T], vals: *Ts) -> T:
    return fn(*vals)

which gives this error:

tst_apply.py:7: error: invalid syntax  [syntax]
    def apply(fn: Callable[[*Ts], T], vals: *Ts) -> T:
                                             ^
Found 1 error in 1 file (errors prevented further checking)

4. using Tuple[*Ts] for TypeVarTuple:

def apply(fn: Callable[[*Ts], T], vals: Tuple[*Ts]) -> T:
    return fn(*vals)

which gives this error:

tst_apply.py:7: error: invalid syntax. Perhaps you forgot a comma?  [syntax]
    def apply(fn: Callable[[*Ts], T], vals: Tuple[*Ts]) -> T:
                                             ^
Found 1 error in 1 file (errors prevented further checking)

Question:

Is there a way to annotate this so the connection between fn parameters and vals is not lost?


Solution

  • In this situation you can use a TypeVarTuple

    Code sample in pyright playground

    from typing import Callable, TypeVar, Tuple, TypeVarTuple
    from typing_extensions import TypeVarTuple, Unpack # < 3.11
    
    T = TypeVar("T")
    Ts = TypeVarTuple('Ts')
    
    # python 3.11
    # Invalid Syntax before 3.11 use Unpack[Ts] for *Ts
    def apply(fn: Callable[[*Ts], T], *vals: *Ts) -> T:
        return fn(*vals)
    
    # < python3.10 and with one argument
    def apply3_10(fn: Callable[[*Ts], T], vals: Tuple[Unpack[Ts]]) -> T:
        return fn(*vals)
    
    def foo(a: int):
        ...
    
    apply3_10(foo, 2)  # Error, not a tuple
    apply3_10(foo, (2,))  # OK
    apply3_10(foo, ("hello",))  # Error
    apply3_10(foo, (2, 3))  # Error
    

    Note, you can only use ParamSpec when you use both P.args and P.kwargs together as you need both to fully form the type-signature. As long as you only use positional arguments to your apply function you are fine with a TypeVarTuple.