pythonpython-typingpython-itertoolsmypy

Type hint for itertools.product doesn't know length of elements


I am adding type hints to an old code of mine. However, I find with a "problem" I don't know how to solve. I have a function that looks like:

def f(x: tuple[tuple[int, int], ...]):
    ...

That is, it accepts one argument which is a tuple of int pairs. At some point in the code I call this function passing as argument the output of itertools.product

from itertools import product

x = tuple(product(range(10), repeat=2))
f(x)

The code works without any problem, since x is a tuple of 100 integer pairs. However, mypy complains with the error

error: Argument 1 to "tuple" has incompatible type "product[tuple[int, ...]]"; expected "Iterable[tuple[int, int]]"  [arg-type]

Is there a way to overcome this problem?


Solution

  • This happens because of the following overload (from typeshed):

    class product(Generic[_T_co]):
        ...
        @overload
        def __new__(cls, *iterables: Iterable[_T1], repeat: int = 1) -> product[tuple[_T1, ...]]: ...
    

    Due to limitations of the type system, the only way to get better results is to have a few more specific overloads at the same place:

    (playgrounds: Mypy, Pyright)

    @overload
    def __new__(cls, iterable: Iterable[_T1], /, *, repeat: Literal[1] = 1) -> product[tuple[_T1]]: ...
    @overload
    def __new__(cls, iterable: Iterable[_T1], /, *, repeat: Literal[2]) -> product[tuple[_T1, _T1]]: ...
    @overload
    def __new__(cls, iterable: Iterable[_T1], /, *, repeat: Literal[3]) -> product[tuple[_T1, _T1, _T1]]: ...
    # And so on
    @overload
    def __new__(cls, iterable: Iterable[_T1], /, *, repeat: int) -> product[tuple[_T1, ...]]: ...
    
    reveal_type(tuple(product(range(10), repeat=2)))  # tuple[tuple[int, int], ...]
    

    However, there aren't any such overloads. Needless to say, this approach doesn't scale very well, since the number of necessary overloads will grow exponentially as the number of iterables increases, though you could try opening an issue to request for it.

    In the meanwhile, it is necessary to use either cast() or # type: ignore:

    f(cast(Iterable[tuple[int, int]], x)
    f(x)  # type: ignore