Overly broad type inference when using zip(*x) to tranpose 2d iterable

[anabelle2001]

Description:

When transposing a 2d iterable x: Iterable[Iterable[T]] using zip(*x), I've noticed typeshed's type annotations result in the return type being inferred as zip[tuple[Any, ...]] rather than the expected zip[tuple[T, ...]]. It would be nice if builtins.py¹ had an overload for this use case.

Expected behavior

When x has type Iterable[Iterable[T]], the expression zip(*x) should be inferred as zip[tuple[T, ...]], preserving the element type information through the transpose operation.

Actual behavior

The expression zip(*x) is inferred as zip[tuple[Any, ...]].

Possible solution

Add the following overloads to zip:

Pre 3.10:

        @overload
        def __new__(
            cls,
            *iterables: Iterable[_T1],
        ) -> zip[tuple[_T1, ...]]: ...

3.10 or later:

        @overload
        def __new__(
            cls,
            *iterables: Iterable[_T1],
            *,
            strict: bool = ...
        ) -> zip[tuple[_T1, ...]]: ...

While I have seen zip(*x) in the wild, I have not seen zip(a,b,c,*x) in any environment before, so I don't think it's necessary to write an overload for that edge case.

Footnotes

1. https://github.com/python/typeshed/blob/494a5d1b98b3522173dd7e0f00f14a32be00456b/stdlib/builtins.pyi#L1813 ↩

[Daverball]

Unfortunately your proposed solution doesn't seem to work currently in mypy.

If we add your overload before the final overload then you get the behavior you want for zip(*x), but then zip(a, b, c, d, e, f, g) would no longer return zip[tuple[Any, ...] and if we do make it the final overload, then zip(*x) will still return zip[tuple[Any, ...]]. There's no way to get one behavior with the spread while getting another with the 7+ parameter case, at least not in mypy. Things might look slightly different in pyright.

At least adding it as the final overload doesn't appear to cause any problems in mypy (but you also don't gain anything), so if you get the correct result in pyright that way, it might still be worth changing it.

It's worth pointing out that there have been some recent discussion about solidifying the behavior of overloads in the typing spec, which also includes the special case of spreading homogenous iterables of unknown length into arguments. So depending on which way these decisions go your solution may end up working better or worse in the future.