improve type for zip_longest

[JelleZijlstra]

Our stub for zip_longest just returned Any for everything. This PR creates a
more precise stub inspired by that for zip(). I had to add additional
overloads to account for the fillvalue parameter.

[JelleZijlstra]

Test case:

from itertools import zip_longest
from typing_extensions import assert_type

assert_type(list(zip_longest([1])), list[tuple[int]])
assert_type(list(zip_longest([1], fillvalue="x")), list[tuple[int]])

assert_type(list(zip_longest([1], ["x"])), list[tuple[int | None, str | None]])
assert_type(list(zip_longest([1], ["x"], fillvalue=1.5)), list[tuple[int | float, str | float]])

assert_type(list(zip_longest([1], ["x"], [1])), list[tuple[int | None, str | None, int | None]])
assert_type(list(zip_longest([1], ["x"], [1], fillvalue=1.5)), list[tuple[int | float, str | float, int | float]])

assert_type(list(zip_longest([1], ["x"], [1], ["x"])), list[tuple[int | None, str | None, int | None, str | None]])
assert_type(
    list(zip_longest([1], ["x"], [1], ["x"], fillvalue=1.5)), list[tuple[int | float, str | float, int | float, str | float]]
)

assert_type(
    list(zip_longest([1], ["x"], [1], ["x"], [1])), list[tuple[int | None, str | None, int | None, str | None, int | None]]
)
assert_type(
    list(zip_longest([1], ["x"], [1], ["x"], [1], fillvalue=1.5)),
    list[tuple[int | float, str | float, int | float, str | float, int | float]],
)

assert_type(list(zip_longest([1], ["x"], [1], ["x"], [1], ["x"])), list[tuple[int | str | None, ...]])
assert_type(list(zip_longest([1], ["x"], [1], ["x"], [1], ["x"], fillvalue=1.5)), list[tuple[int | str | float, ...]])

This passes pyright, but mypy infers list[Any] for all of them, probably because of a type context bug. I'll investigate that further.

[AlexWaygood]

If we take the rich hit as an example, it seems a little unfortunate:

for index, (word, next_word) in enumerate(
    zip_longest(words, words[1:])
):
    tokens.append(word)

Mypy thinks that word could be None here, because fillvalue hasn't been specified, so it assumes the Iterator will be yielding (str | None, str | None) tuples. But in actual fact, we know that word can only ever be str, as word will always be taken from the longer iterable.

This seems like pretty idiomatic usage of zip_longest :/

[JelleZijlstra]

With python/mypy#12612 instead mypy passes my tests except for the last one, where it infers List[Tuple[object, ...]] because it uses a join instead of a union.

Looking at the mypy-primer output:

• https://github.com/konradhalas/dacite/blob/61182ad8b8afc2ac74f51bbf2a4fb7f45d35da74/dacite/core.py#L137 I think this is a true positive, it's not obvious from the code which iterable is longer.
• https://github.com/Textualize/rich/blob/3b36864ad00ff10276bc86e089d4a08f6f34b890/rich/containers.py#L159 does zip_longest(words, words[1:]), so obviously the first entry will not be None. We could handle this by returning _T | Any instead of _T | None in the overloads without fillvalue.
• https://github.com/pydata/xarray/blob/586992e8d2998751cb97b1cab4d3caa9dca116e0/xarray/core/formatting.py#L199 I think shouldn't be using zip_longest.
• https://github.com/kornia/kornia/blob/42c5a88453d5b79054c2a0b5c826b4586ba505bb/kornia/augmentation/container/image.py#L329 looks unsafe
• https://github.com/kornia/kornia/blob/42c5a88453d5b79054c2a0b5c826b4586ba505bb/kornia/augmentation/container/augment.py#L219 is similar
• https://github.com/schemathesis/schemathesis/blob/2cca06135648c26262f32dd5664d2851de3121fb/src/schemathesis/models.py#L542 is like the rich case, it does zip_longest(url, url[1:]).

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

kornia (https://github.com/kornia/kornia)
+ kornia/augmentation/container/image.py:331: error: Invalid index type "Union[str, Any]" for "List[ParamItem]"; expected type "SupportsIndex"  [index]
+ kornia/augmentation/container/image.py:340: error: Item "None" of "Union[ParamItem, Any, None]" has no attribute "data"  [union-attr]
+ kornia/augmentation/container/image.py:340: error: Argument 2 to "inverse" of "ImageSequential" has incompatible type "Union[Dict[Any, Any], List[Any], None, Any]"; expected "Optional[List[ParamItem]]"  [arg-type]
+ kornia/augmentation/container/augment.py:221: error: Invalid index type "Union[str, Any]" for "List[ParamItem]"; expected type "SupportsIndex"  [index]