Typing for callable inputs in itertools.starmap

[aakhavanQC]

The type annotations for itertools.starmap allow any iterables for any function arguments: https://github.com/python/typeshed/blob/11c7821a79a8ab7e1982f3ab506db16f1c4a22a9/stdlib/itertools.pyi#L112-L116

This means I can do something like:

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

itertools.starmap(myfunc, [["foo, "bar"]])

and type checkers won't complain.

map gets around this to some extent with overloads for callables with up to 5 arguments (but if you need 6 or more, you're out of luck!) https://github.com/python/typeshed/blob/11c7821a79a8ab7e1982f3ab506db16f1c4a22a9/stdlib/builtins.pyi#L1629-L1670

I couldn't figure out a way to ge that to work with starmap, though, because as far as I can tell the python type system doesn't have any way of specifying an iterable that returns a particular sequence of types. The best you can do is something like def __new__(cls, function: Callable[[_T1, _T2], T], iterable: Iterable[Iterable[_T1 | _T2]], /), which is still better in my opinion (it prevents invalid argument types), but it doesn't stop you from swapping function arguments.

Is there a way to do this? Or does it just need a new language feature? It would be great if we could do something like

P = ParamSpec("P")
def __new__(cls, function: Callable[P, T], iterable: Iterable[P.args], /)

[smelnikov]

As far as I know, this type can’t be fully expressed. The closest approach is using Unpack with tuple, which allows specific types to be bound to specific positions.

_Ts = TypeVarTuple("_Ts")

@disjoint_base
class starmap(Generic[_T_co,]):
    def __new__(cls, function: Callable[[Unpack[Ts]], _T], iterable: Iterable[tuple[Unpack[Ts]]], /) -> "starmap[_T]": ...
    # ...

But it only works with iterables of tuples. Also, this syntax is supported only in Python 3.11+, if I’m not mistaken.

def args(a: int, b: str) -> int: ...
def varargs(*args: int) -> int: ...

starmap(args, [(1, "a"), (3, "b")]) # ok
starmap(varargs, [(1, 2, 3), (4, 5, 6)]) # ok

starmap(args, [[1, "a"], [3, "b"]]) # list[int | str] is not assignable to tuple[int, str]
starmap(varargs, [[1, 2, 3], [4, 5, 6]]) # list[int] is not assignable to tuple[int, ...]

If you need a bit more type safety, you can create a custom starmap wrapper to handle this specific case in your code.