diff --git a/pytype/stubs/builtins/builtins.pytd b/pytype/stubs/builtins/builtins.pytd index 8effeca99..aa1a4e41d 100644 --- a/pytype/stubs/builtins/builtins.pytd +++ b/pytype/stubs/builtins/builtins.pytd @@ -172,21 +172,37 @@ def sum(iterable: Iterable[_T], start: _T) -> _T: ... def vars(*args, **kwargs) -> dict[str, Any]: ... if sys.version_info >= (3, 10): - def zip(*, strict: bool = ...) -> Iterator[nothing]: ... - def zip(seq1: Iterable[nothing], *, strict: bool = ...) -> Iterator[nothing]: ... - def zip(seq1: Iterable[_T], *, strict: bool = ...) -> Iterator[Tuple[_T]]: ... - def zip(seq1: Iterable[nothing], seq2: Iterable, *, strict: bool = ...) -> Iterator[nothing]: ... - def zip(seq1: Iterable, seq2: Iterable[nothing], *, strict: bool = ...) -> Iterator[nothing]: ... - def zip(seq1: Iterable[_T], seq2: Iterable[_T2], *, strict: bool = ...) -> Iterator[Tuple[_T, _T2]]: ... - def zip(seq1, seq2, seq3, *seqs: Iterable, strict: bool = ...) -> Iterator[tuple]: ... + class zip(Iterable[_S]): + @overload + def __new__(cls, *, strict: bool = ...) -> zip[nothing]: ... + @overload + def __new__(cls, seq: Iterable[nothing], *, strict: bool = ...) -> zip[nothing]: ... + @overload + def __new__(cls, seq: Iterable[_T], *, strict: bool = ...) -> zip[Tuple[_T]]: ... + @overload + def __new__(cls, seq1: Iterable[nothing], seq2: Iterable, *, strict: bool = ...) -> zip[nothing]: ... + @overload + def __new__(cls, seq1: Iterable, seq2: Iterable[nothing], *, strict: bool = ...) -> zip[nothing]: ... + @overload + def __new__(cls, seq1: Iterable[_T], seq2: Iterable[_T2], *, strict: bool = ...) -> zip[Tuple[_T, _T2]]: ... + @overload + def __new__(cls, seq1, seq2, seq3, *seqs: Iterable, strict: bool = ...) -> zip[tuple]: ... else: - def zip() -> Iterator[nothing]: ... - def zip(seq1: Iterable[nothing]) -> Iterator[nothing]: ... - def zip(seq1: Iterable[_T]) -> Iterator[Tuple[_T]]: ... - def zip(seq1: Iterable[nothing], seq2: Iterable) -> Iterator[nothing]: ... - def zip(seq1: Iterable, seq2: Iterable[nothing]) -> Iterator[nothing]: ... - def zip(seq1: Iterable[_T], seq2: Iterable[_T2]) -> Iterator[Tuple[_T, _T2]]: ... - def zip(seq1, seq2, seq3, *seqs: Iterable) -> Iterator[tuple]: ... + class zip(Iterable[_S]): + @overload + def __new__(cls) -> zip[nothing]: ... + @overload + def __new__(cls, seq: Iterable[nothing]) -> zip[nothing]: ... + @overload + def __new__(cls, seq: Iterable[_T]) -> zip[Tuple[_T]]: ... + @overload + def __new__(cls, seq1: Iterable[nothing], seq2: Iterable) -> zip[nothing]: ... + @overload + def __new__(cls, seq1: Iterable, seq2: Iterable[nothing]) -> zip[nothing]: ... + @overload + def __new__(cls, seq1: Iterable[_T], seq2: Iterable[_T2]) -> zip[Tuple[_T, _T2]]: ... + @overload + def __new__(cls, seq1, seq2, seq3, *seqs: Iterable) -> zip[tuple]: ... # Keyword constants cannot be assigned here; they are defined in the parser. # False: bool = ... diff --git a/pytype/test_data/complex.pyi b/pytype/test_data/complex.pyi index 85e63a3a4..69ed38bd7 100644 --- a/pytype/test_data/complex.pyi +++ b/pytype/test_data/complex.pyi @@ -1,7 +1,7 @@ -from typing import Any, Iterator +from typing import Any v = ... # type: list w = ... # type: list x = ... # type: list y = ... # type: list -z = ... # type: Iterator[tuple] +z = ... # type: zip[tuple] diff --git a/pytype/tests/test_builtins4.py b/pytype/tests/test_builtins4.py index 78ba0c097..6ca7e4692 100644 --- a/pytype/tests/test_builtins4.py +++ b/pytype/tests/test_builtins4.py @@ -451,11 +451,11 @@ def test_zip(self): """) self.assertTypesMatchPytd(ty, """ from typing import Iterator, Tuple, Union - a : Iterator[nothing] - b : Iterator[Tuple[Union[int, complex]]] - c : Iterator[nothing] - d : Iterator[nothing] - e : Iterator[Tuple[complex, int]] + a: zip[nothing] + b: zip[Tuple[Union[int, complex]]] + c: zip[nothing] + d: zip[nothing] + e: zip[Tuple[complex, int]] """) def test_dict(self): @@ -615,7 +615,7 @@ def test_iterator_builtins(self): self.assertTypesMatchPytd(ty, """ from typing import Iterator, Tuple v1 : Iterator[int] - v2 : Iterator[Tuple[int, int]] + v2 : zip[Tuple[int, int]] """) def test_next(self):