From d80cc7ab23c9e427e9c13104bf45d52881cabfd0 Mon Sep 17 00:00:00 2001 From: Jelle Zijlstra Date: Sun, 17 Apr 2022 14:55:12 -0700 Subject: [PATCH 1/3] improve type for zip_longest --- stdlib/itertools.pyi | 96 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 84 insertions(+), 12 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index 520d8ecb2c37..e03a93397e91 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -1,5 +1,5 @@ import sys -from _typeshed import Self, _T_co +from _typeshed import Self from typing import Any, Callable, Generic, Iterable, Iterator, SupportsComplex, SupportsFloat, SupportsInt, TypeVar, overload from typing_extensions import Literal, SupportsIndex, TypeAlias @@ -9,6 +9,14 @@ if sys.version_info >= (3, 9): _T = TypeVar("_T") _S = TypeVar("_S") _N = TypeVar("_N", int, float, SupportsFloat, SupportsInt, SupportsIndex, SupportsComplex) +_T_co = TypeVar("_T_co") +_T1 = TypeVar("_T1") +_T2 = TypeVar("_T2") +_T3 = TypeVar("_T3") +_T4 = TypeVar("_T4") +_T5 = TypeVar("_T5") +_T6 = TypeVar("_T6") + _Step: TypeAlias = int | float | SupportsFloat | SupportsInt | SupportsIndex | SupportsComplex Predicate: TypeAlias = Callable[[_T], object] @@ -76,9 +84,6 @@ class filterfalse(Iterator[_T], Generic[_T]): def __iter__(self: Self) -> Self: ... def __next__(self) -> _T: ... -_T1 = TypeVar("_T1") -_T2 = TypeVar("_T2") - class groupby(Iterator[tuple[_T, Iterator[_S]]], Generic[_T, _S]): @overload def __new__(cls, iterable: Iterable[_T1], key: None = ...) -> groupby[_T1, _T1]: ... @@ -107,15 +112,82 @@ class takewhile(Iterator[_T], Generic[_T]): def tee(__iterable: Iterable[_T], __n: int = ...) -> tuple[Iterator[_T], ...]: ... -class zip_longest(Iterator[Any]): - def __init__(self, *p: Iterable[Any], fillvalue: Any = ...) -> None: ... +class zip_longest(Iterator[_T_co], Generic[_T_co]): + # one iterable (fillvalue doesn't matter) + @overload + def __new__(cls, __iter1: Iterable[_T1], *, fillvalue: object = ...) -> zip_longest[tuple[_T1]]: ... + # two iterables + @overload + def __new__(cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2]) -> zip_longest[tuple[_T1 | None, _T2 | None]]: ... + @overload + def __new__( + cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], *, fillvalue: _T + ) -> zip_longest[tuple[_T1 | _T, _T2 | _T]]: ... + # three iterables + @overload + def __new__( + cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3] + ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None]]: ... + @overload + def __new__( + cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], *, fillvalue: _T + ) -> zip_longest[tuple[_T1 | _T, _T2 | _T, _T3 | _T]]: ... + # four iterables + @overload + def __new__( + cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], __iter4: Iterable[_T4] + ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None, _T4 | None]]: ... + @overload + def __new__( + cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], __iter4: Iterable[_T4], *, fillvalue: _T + ) -> zip_longest[tuple[_T1 | _T, _T2 | _T, _T3 | _T, _T4 | _T]]: ... + # five iterables + @overload + def __new__( + cls, + __iter1: Iterable[_T1], + __iter2: Iterable[_T2], + __iter3: Iterable[_T3], + __iter4: Iterable[_T4], + __iter5: Iterable[_T5], + ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None, _T4 | None, _T5 | None]]: ... + @overload + def __new__( + cls, + __iter1: Iterable[_T1], + __iter2: Iterable[_T2], + __iter3: Iterable[_T3], + __iter4: Iterable[_T4], + __iter5: Iterable[_T5], + *, + fillvalue: _T, + ) -> zip_longest[tuple[_T1 | _T, _T2 | _T, _T3 | _T, _T4 | _T, _T5 | _T]]: ... + # six or more iterables + @overload + def __new__( + cls, + __iter1: Iterable[_T], + __iter2: Iterable[_T], + __iter3: Iterable[_T], + __iter4: Iterable[_T], + __iter5: Iterable[_T], + __iter6: Iterable[_T], + *iterables: Iterable[_T], + ) -> zip_longest[tuple[_T | None, ...]]: ... + @overload + def __new__( + cls, + __iter1: Iterable[_T], + __iter2: Iterable[_T], + __iter3: Iterable[_T], + __iter4: Iterable[_T], + __iter5: Iterable[_T], + __iter6: Iterable[_T], + *iterables: Iterable[_T], + fillvalue: _T, + ) -> zip_longest[tuple[_T, ...]]: ... def __iter__(self: Self) -> Self: ... - def __next__(self) -> Any: ... - -_T3 = TypeVar("_T3") -_T4 = TypeVar("_T4") -_T5 = TypeVar("_T5") -_T6 = TypeVar("_T6") + def __next__(self) -> _T_co: ... class product(Iterator[_T_co], Generic[_T_co]): @overload From 5326b226c0eb383ba21d2efb4f9450d5ccb2bbb6 Mon Sep 17 00:00:00 2001 From: Jelle Zijlstra Date: Sun, 17 Apr 2022 15:01:10 -0700 Subject: [PATCH 2/3] make _T_co actually covariant --- stdlib/itertools.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index e03a93397e91..1427d6f2f8e0 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -9,7 +9,7 @@ if sys.version_info >= (3, 9): _T = TypeVar("_T") _S = TypeVar("_S") _N = TypeVar("_N", int, float, SupportsFloat, SupportsInt, SupportsIndex, SupportsComplex) -_T_co = TypeVar("_T_co") +_T_co = TypeVar("_T_co", covariant=True) _T1 = TypeVar("_T1") _T2 = TypeVar("_T2") _T3 = TypeVar("_T3") From 61619abcbfb717d86914bbe9fa331efc04fcaa3b Mon Sep 17 00:00:00 2001 From: Jelle Zijlstra Date: Sun, 17 Apr 2022 15:27:06 -0700 Subject: [PATCH 3/3] fall back to Any --- stdlib/itertools.pyi | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index 1427d6f2f8e0..35884bc30ac0 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -118,7 +118,10 @@ class zip_longest(Iterator[_T_co], Generic[_T_co]): def __new__(cls, __iter1: Iterable[_T1], *, fillvalue: object = ...) -> zip_longest[tuple[_T1]]: ... # two iterables @overload - def __new__(cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2]) -> zip_longest[tuple[_T1 | None, _T2 | None]]: ... + # In the overloads without fillvalue, all of the tuple members could theoretically be None, + # but we return Any instead to avoid false positives for code where we know one of the iterables + # is longer. + def __new__(cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2]) -> zip_longest[tuple[_T1 | Any, _T2 | Any]]: ... @overload def __new__( cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], *, fillvalue: _T @@ -127,7 +130,7 @@ class zip_longest(Iterator[_T_co], Generic[_T_co]): @overload def __new__( cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3] - ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None]]: ... + ) -> zip_longest[tuple[_T1 | Any, _T2 | Any, _T3 | Any]]: ... @overload def __new__( cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], *, fillvalue: _T @@ -136,7 +139,7 @@ class zip_longest(Iterator[_T_co], Generic[_T_co]): @overload def __new__( cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], __iter4: Iterable[_T4] - ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None, _T4 | None]]: ... + ) -> zip_longest[tuple[_T1 | Any, _T2 | Any, _T3 | Any, _T4 | Any]]: ... @overload def __new__( cls, __iter1: Iterable[_T1], __iter2: Iterable[_T2], __iter3: Iterable[_T3], __iter4: Iterable[_T4], *, fillvalue: _T @@ -150,7 +153,7 @@ class zip_longest(Iterator[_T_co], Generic[_T_co]): __iter3: Iterable[_T3], __iter4: Iterable[_T4], __iter5: Iterable[_T5], - ) -> zip_longest[tuple[_T1 | None, _T2 | None, _T3 | None, _T4 | None, _T5 | None]]: ... + ) -> zip_longest[tuple[_T1 | Any, _T2 | Any, _T3 | Any, _T4 | Any, _T5 | Any]]: ... @overload def __new__( cls, @@ -173,7 +176,7 @@ class zip_longest(Iterator[_T_co], Generic[_T_co]): __iter5: Iterable[_T], __iter6: Iterable[_T], *iterables: Iterable[_T], - ) -> zip_longest[tuple[_T | None, ...]]: ... + ) -> zip_longest[tuple[_T | Any, ...]]: ... @overload def __new__( cls,