Convert itertools functions to classes (#5211)

stdlib/itertools.pyi

@@ -1,60 +1,123 @@
 import sys
-from typing import Any, Callable, Generic, Iterable, Iterator, Optional, Tuple, TypeVar, overload
-from typing_extensions import Literal
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Iterator,
+    Optional,
+    SupportsComplex,
+    SupportsFloat,
+    SupportsInt,
+    Tuple,
+    Type,
+    TypeVar,
+    overload,
+)
+from typing_extensions import Literal, SupportsIndex
 
 _T = TypeVar("_T")
 _S = TypeVar("_S")
-_N = TypeVar("_N", int, float)
+_N = TypeVar("_N", int, float, SupportsFloat, SupportsInt, SupportsIndex, SupportsComplex)
+_NStep = TypeVar("_NStep", int, float, SupportsFloat, SupportsInt, SupportsIndex, SupportsComplex)
+
 Predicate = Callable[[_T], object]
 
-def count(start: _N = ..., step: _N = ...) -> Iterator[_N]: ...  # more general types?
+# Technically count can take anything that implements a number protocol and has an add method
+# but we can't enforce the add method
+class count(Iterator[_N], Generic[_N]):
+    @overload
+    def __new__(cls) -> count[int]: ...
+    @overload
+    def __new__(cls, start: _N, step: _NStep = ...) -> count[_N]: ...
+    def __next__(self) -> _N: ...
+    def __iter__(self) -> Iterator[_N]: ...
 
 class cycle(Iterator[_T], Generic[_T]):
     def __init__(self, __iterable: Iterable[_T]) -> None: ...
     def __next__(self) -> _T: ...
     def __iter__(self) -> Iterator[_T]: ...
 
-@overload
-def repeat(object: _T) -> Iterator[_T]: ...
-@overload
-def repeat(object: _T, times: int) -> Iterator[_T]: ...
-
-if sys.version_info >= (3, 8):
+class repeat(Iterator[_T], Generic[_T]):
     @overload
-    def accumulate(iterable: Iterable[_T], func: None = ..., *, initial: Optional[_T] = ...) -> Iterator[_T]: ...
+    def __init__(self, object: _T) -> None: ...
     @overload
-    def accumulate(iterable: Iterable[_T], func: Callable[[_S, _T], _S], *, initial: Optional[_S] = ...) -> Iterator[_S]: ...
+    def __init__(self, object: _T, times: int) -> None: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Iterator[_T]: ...
 
-else:
-    def accumulate(iterable: Iterable[_T], func: Optional[Callable[[_T, _T], _T]] = ...) -> Iterator[_T]: ...
+class accumulate(Iterator[_T], Generic[_T]):
+    if sys.version_info >= (3, 8):
+        @overload
+        def __init__(self, iterable: Iterable[_T], func: None = ..., *, initial: Optional[_T] = ...) -> None: ...
+        @overload
+        def __init__(self, iterable: Iterable[_S], func: Callable[[_T, _S], _T], *, initial: Optional[_T] = ...) -> None: ...
+    else:
+        def __init__(self, iterable: Iterable[_T], func: Optional[Callable[[_T, _T], _T]] = ...) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _T: ...
 
 class chain(Iterator[_T], Generic[_T]):
     def __init__(self, *iterables: Iterable[_T]) -> None: ...
     def __next__(self) -> _T: ...
     def __iter__(self) -> Iterator[_T]: ...
-    @staticmethod
-    def from_iterable(__iterable: Iterable[Iterable[_S]]) -> Iterator[_S]: ...
+    @classmethod
+    # We use Type and not Type[_S] to not lose the type inference from __iterable
+    def from_iterable(cls: Type[Any], __iterable: Iterable[Iterable[_S]]) -> Iterator[_S]: ...
 
-def compress(data: Iterable[_T], selectors: Iterable[Any]) -> Iterator[_T]: ...
-def dropwhile(__predicate: Predicate[_T], __iterable: Iterable[_T]) -> Iterator[_T]: ...
-def filterfalse(__predicate: Optional[Predicate[_T]], __iterable: Iterable[_T]) -> Iterator[_T]: ...
-@overload
-def groupby(iterable: Iterable[_T], key: None = ...) -> Iterator[Tuple[_T, Iterator[_T]]]: ...
-@overload
-def groupby(iterable: Iterable[_T], key: Callable[[_T], _S]) -> Iterator[Tuple[_S, Iterator[_T]]]: ...
-@overload
-def islice(__iterable: Iterable[_T], __stop: Optional[int]) -> Iterator[_T]: ...
-@overload
-def islice(
-    __iterable: Iterable[_T], __start: Optional[int], __stop: Optional[int], __step: Optional[int] = ...
-) -> Iterator[_T]: ...
-def starmap(__function: Callable[..., _S], __iterable: Iterable[Iterable[Any]]) -> Iterator[_S]: ...
-def takewhile(__predicate: Predicate[_T], __iterable: Iterable[_T]) -> Iterator[_T]: ...
-def tee(__iterable: Iterable[_T], __n: int = ...) -> Tuple[Iterator[_T], ...]: ...
-def zip_longest(*p: Iterable[Any], fillvalue: Any = ...) -> Iterator[Any]: ...
+class compress(Iterator[_T], Generic[_T]):
+    def __init__(self, data: Iterable[_T], selectors: Iterable[Any]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _T: ...
+
+class dropwhile(Iterator[_T], Generic[_T]):
+    def __init__(self, __predicate: Predicate[_T], __iterable: Iterable[_T]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _T: ...
+
+class filterfalse(Iterator[_T], Generic[_T]):
+    def __init__(self, __predicate: Optional[Predicate[_T]], __iterable: Iterable[_T]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    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]: ...
+    @overload
+    def __new__(cls, iterable: Iterable[_T1], key: Callable[[_T1], _T2]) -> groupby[_T2, _T1]: ...
+    def __iter__(self) -> Iterator[Tuple[_T, Iterator[_S]]]: ...
+    def __next__(self) -> Tuple[_T, Iterator[_S]]: ...
+
+class islice(Iterator[_T], Generic[_T]):
+    @overload
+    def __init__(self, __iterable: Iterable[_T], __stop: Optional[int]) -> None: ...
+    @overload
+    def __init__(
+        self, __iterable: Iterable[_T], __start: Optional[int], __stop: Optional[int], __step: Optional[int] = ...
+    ) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _T: ...
+
+class starmap(Iterator[_T], Generic[_T]):
+    def __init__(self, __function: Callable[..., _T], __iterable: Iterable[Iterable[Any]]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _T: ...
+
+class takewhile(Iterator[_T], Generic[_T]):
+    def __init__(self, __predicate: Predicate[_T], __iterable: Iterable[_T]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __next__(self) -> _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: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __next__(self) -> Any: ...
+
 _T3 = TypeVar("_T3")
 _T4 = TypeVar("_T4")
 _T5 = TypeVar("_T5")
@@ -97,7 +160,12 @@ def product(
 def product(*iterables: Iterable[_T1], repeat: int) -> Iterator[Tuple[_T1, ...]]: ...
 @overload
 def product(*iterables: Iterable[Any], repeat: int = ...) -> Iterator[Tuple[Any, ...]]: ...
-def permutations(iterable: Iterable[_T], r: Optional[int] = ...) -> Iterator[Tuple[_T, ...]]: ...
+
+class permutations(Iterator[Tuple[_T, ...]], Generic[_T]):
+    def __init__(self, iterable: Iterable[_T], r: Optional[int] = ...) -> None: ...
+    def __iter__(self) -> Iterator[Tuple[_T, ...]]: ...
+    def __next__(self) -> Tuple[_T, ...]: ...
+
 @overload
 def combinations(iterable: Iterable[_T], r: Literal[2]) -> Iterator[Tuple[_T, _T]]: ...
 @overload
@@ -108,4 +176,8 @@ def combinations(iterable: Iterable[_T], r: Literal[4]) -> Iterator[Tuple[_T, _T
 def combinations(iterable: Iterable[_T], r: Literal[5]) -> Iterator[Tuple[_T, _T, _T, _T, _T]]: ...
 @overload
 def combinations(iterable: Iterable[_T], r: int) -> Iterator[Tuple[_T, ...]]: ...
-def combinations_with_replacement(iterable: Iterable[_T], r: int) -> Iterator[Tuple[_T, ...]]: ...
+
+class combinations_with_replacement(Iterator[Tuple[_T, ...]], Generic[_T]):
+    def __init__(self, iterable: Iterable[_T], r: int) -> None: ...
+    def __iter__(self) -> Iterator[Tuple[_T, ...]]: ...
+    def __next__(self) -> Tuple[_T, ...]: ...

tests/stubtest_whitelists/py36.txt

@@ -28,7 +28,6 @@ importlib.metadata  # Added in 3.8
 importlib.resources  # Added in 3.7
 io.StringIO.readline
 ipaddress._BaseNetwork.__init__
-itertools.accumulate
 json.loads
 mmap.ACCESS_DEFAULT
 multiprocessing.shared_memory

tests/stubtest_whitelists/py37.txt

@@ -38,7 +38,6 @@ http.client.HTTPSConnection.__init__
 http.server.SimpleHTTPRequestHandler.__init__
 importlib.metadata  # Added in 3.8
 ipaddress._BaseNetwork.__init__
-itertools.accumulate
 json.loads
 macurl2path  # removed in 3.7
 multiprocessing.shared_memory

tests/stubtest_whitelists/py3_common.txt

@@ -199,22 +199,8 @@ ipaddress._BaseAddress.is_unspecified
 ipaddress._BaseAddress.max_prefixlen
 ipaddress._BaseAddress.packed
 ipaddress._BaseNetwork.max_prefixlen
-itertools.accumulate  # not a function at runtime
-itertools.chain.from_iterable
-itertools.combinations  # not a function at runtime
-itertools.combinations_with_replacement
-itertools.compress
-itertools.count
-itertools.dropwhile
-itertools.filterfalse
-itertools.groupby  # not a function at runtime
-itertools.islice  # not a function at runtime
-itertools.permutations
-itertools.product  # not a function at runtime
-itertools.repeat  # not a function at runtime
-itertools.starmap
-itertools.takewhile
-itertools.zip_longest
+itertools.combinations  # not a function at runtime, cannot represent with __init__ or __new__
+itertools.product  # not a function at runtime, cannot represent with __init__ or __new__
 lib2to3.pygram.pattern_symbols
 lib2to3.pygram.python_symbols
 lib2to3.pytree.Base.__new__