Typing for callable inputs in itertools.starmap

[aakhavanQC]

The type annotations for itertools.starmap allow any iterables for any function arguments:

typeshed/stdlib/itertools.pyi

Lines 112 to 116 in 11c7821

	@disjoint_base
	class starmap(Generic[_T_co]):
	def __new__(cls, function: Callable[..., _T], iterable: Iterable[Iterable[Any]], /) -> starmap[_T]: ...
	def __iter__(self) -> Self: ...
	def __next__(self) -> _T_co: ...

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!)

typeshed/stdlib/builtins.pyi

Lines 1629 to 1670 in 11c7821

	@overload
	def __new__(cls, func: Callable[[_T1], _S], iterable: Iterable[_T1], /) -> Self: ...
	@overload
	def __new__(cls, func: Callable[[_T1, _T2], _S], iterable: Iterable[_T1], iter2: Iterable[_T2], /) -> Self: ...
	@overload
	def __new__(
	cls, func: Callable[[_T1, _T2, _T3], _S], iterable: Iterable[_T1], iter2: Iterable[_T2], iter3: Iterable[_T3], /
	) -> Self: ...
	@overload
	def __new__(
	cls,
	func: Callable[[_T1, _T2, _T3, _T4], _S],
	iterable: Iterable[_T1],
	iter2: Iterable[_T2],
	iter3: Iterable[_T3],
	iter4: Iterable[_T4],
	/,
	) -> Self: ...
	@overload
	def __new__(
	cls,
	func: Callable[[_T1, _T2, _T3, _T4, _T5], _S],
	iterable: Iterable[_T1],
	iter2: Iterable[_T2],
	iter3: Iterable[_T3],
	iter4: Iterable[_T4],
	iter5: Iterable[_T5],
	/,
	) -> Self: ...
	@overload
	def __new__(
	cls,
	func: Callable[..., _S],
	iterable: Iterable[Any],
	iter2: Iterable[Any],
	iter3: Iterable[Any],
	iter4: Iterable[Any],
	iter5: Iterable[Any],
	iter6: Iterable[Any],
	/,
	*iterables: Iterable[Any],
	) -> Self: ...

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.