itertools allowlist fixes

python · srittau · Apr 12, 2021 · Mar 28, 2021 · Mar 28, 2021 · Apr 11, 2021
commit 6034334a00d20f59040309b9a6098d387c9eec29
diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi
@@ -1,13 +1,20 @@
 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, Tuple, Type, TypeVar, overload, SupportsFloat, SupportsInt, SupportsComplex
+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, Su
10000
pportsComplex)
+
 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]):
+    def __init__(self, start: _N = ..., step: _NStep = ...) -> None: ...
+    def __next__(self) -> _N: ...
+    def __iter__(self) -> Iterator[_N]: ...
 
 class cycle(Iterator[_T], Generic[_T]):
     def __init__(self, __iterable: Iterable[_T]) -> None: ...
@@ -32,12 +39,25 @@ 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, __iterable: Iterable[Iterable[_S]]) -> Iterator[_S]: ...
+
+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: ...
 
-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
@@ -48,10 +68,23 @@ def islice(__iterable: Iterable[_T], __stop: Optional[int]) -> Iterator[_T]: ...
 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]: ...
+
+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], ...]: ...
-def zip_longest(*p: Iterable[Any], fillvalue: Any = ...) -> Iterator[Any]: ...
+
+class zip_longest(Iterator):
+    def __init__(self, *p: Iterable[Any], fillvalue: Any = ...) -> None: ...
+    def __iter__(self) -> Iterator: ...
+    def __next__(self) -> Any: ...
 
 _T1 = TypeVar("_T1")
 _T2 = TypeVar("_T2")
@@ -97,7 +130,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 +146,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, ...]: ...
+
diff --git a/tests/stubtest_whitelists/py36.txt b/tests/stubtest_whitelists/py36.txt
@@ -30,7 +30,6 @@ importlib.metadata
 importlib.resources
 io.StringIO.readline
 ipaddress._BaseNetwork.__init__
-itertools.accumulate
 json.loads
 mmap.ACCESS_DEFAULT
 multiprocessing.shared_memory

diff --git a/tests/stubtest_whitelists/py37.txt b/tests/stubtest_whitelists/py37.txt
@@ -38,7 +38,6 @@ http.client.HTTPSConnection.__init__
 http.server.SimpleHTTPRequestHandler.__init__
 importlib.metadata
 ipaddress._BaseNetwork.__init__
-itertools.accumulate
 json.loads
 macurl2path  # removed in 3.7
 multiprocessing.shared_memory

diff --git a/tests/stubtest_whitelists/py3_common.txt b/tests/stubtest_whitelists/py3_common.txt
@@ -247,16 +247,6 @@ ipaddress._BaseAddress.is_unspecified
 ipaddress._BaseAddress.max_prefixlen
 ipaddress._BaseAddress.packed
 ipaddress._BaseNetwork.max_prefixlen
-itertools.chain.from_iterable
-itertools.combinations_with_replacement
-itertools.compress
-itertools.count
-itertools.dropwhile
-itertools.filterfalse
-itertools.permutations
-itertools.starmap
-    itertools.takewhile
-itertools.zip_longest
 lib2to3.pygram.pattern_symbols
 lib2to3.pygram.python_symbols
 lib2to3.pytree.Base.__new__