From 6034334a00d20f59040309b9a6098d387c9eec29 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Sun, 28 Mar 2021 03:29:39 +0300 Subject: [PATCH 1/5] itertools allowlist fixes --- stdlib/itertools.pyi | 70 +++++++++++++++++++----- tests/stubtest_whitelists/py36.txt | 1 - tests/stubtest_whitelists/py37.txt | 1 - tests/stubtest_whitelists/py3_common.txt | 10 ---- 4 files changed, 56 insertions(+), 26 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index aef3aac78635..b022a47e9907 100644 --- 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, 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]): + 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 index d76c12b05e30..27b8945aa806 100644 --- 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 index 2d3dbb584ece..84bc3f7bd2b2 100644 --- 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 index f06288ec1d19..a0fd798b88de 100644 --- 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__ From ecc7f2f1fe5898be7d640d49dbaf5065a2066575 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Sun, 28 Mar 2021 03:54:14 +0300 Subject: [PATCH 2/5] Fix accumulate issue and formatting --- stdlib/itertools.pyi | 44 ++++++++++++++++++++++++++++++-------------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index b022a47e9907..e08585b34703 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -1,5 +1,19 @@ import sys -from typing import Any, Callable, Generic, Iterable, Iterator, Optional, Tuple, Type, TypeVar, overload, SupportsFloat, SupportsInt, SupportsComplex +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") @@ -26,14 +40,16 @@ def repeat(object: _T) -> Iterator[_T]: ... @overload def repeat(object: _T, times: int) -> Iterator[_T]: ... -if sys.version_info >= (3, 8): - @overload - def accumulate(iterable: Iterable[_T], func: None = ..., *, initial: Optional[_T] = ...) -> Iterator[_T]: ... - @overload - def accumulate(iterable: Iterable[_T], func: Callable[[_S, _T], _S], *, initial: Optional[_S] = ...) -> Iterator[_S]: ... - -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: ... @@ -41,7 +57,7 @@ class chain(Iterator[_T], Generic[_T]): def __iter__(self) -> Iterator[_T]: ... @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]: ... + def from_iterable(cls: Type[Any], __iterable: Iterable[Iterable[_S]]) -> Iterator[_S]: ... class compress(Iterator[_T], Generic[_T]): def __init__(self, data: Iterable[_T], selectors: Iterable[Any]) -> None: ... @@ -54,7 +70,7 @@ class dropwhile(Iterator[_T], Generic[_T]): def __next__(self) -> _T: ... class filterfalse(Iterator[_T], Generic[_T]): - def __init__(self, __predicate: Optional[Predicate[_T]], __iterable: Iterable[_T]) -> None: ... + def __init__(self, __predicate: Optional[Predicate[_T]], __iterable: Iterable[_T]) -> None: ... def __iter__(self) -> Iterator[_T]: ... def __next__(self) -> _T: ... @@ -81,9 +97,9 @@ class takewhile(Iterator[_T], Generic[_T]): def tee(__iterable: Iterable[_T], __n: int = ...) -> Tuple[Iterator[_T], ...]: ... -class zip_longest(Iterator): +class zip_longest(Iterator[Any]): def __init__(self, *p: Iterable[Any], fillvalue: Any = ...) -> None: ... - def __iter__(self) -> Iterator: ... + def __iter__(self) -> Iterator[Any]: ... def __next__(self) -> Any: ... _T1 = TypeVar("_T1") @@ -146,8 +162,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, ...]]: ... + 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, ...]: ... - From d23b97fb26178bee401de824fdde923263f93987 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Mon, 12 Apr 2021 07:21:46 +0300 Subject: [PATCH 3/5] Fix more functions --- stdlib/itertools.pyi | 27 +++++++++++++++--------- tests/stubtest_whitelists/py3_common.txt | 9 +++----- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index e08585b34703..e6f223e725a4 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -35,10 +35,13 @@ class cycle(Iterator[_T], Generic[_T]): 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]: ... +class repeat(Iterator[_T], Generic[_T]): + @overload + def __init__(self, object: _T) -> None: ... + @overload + def __init__(self, object: _T, times: int) -> None: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Iterator[_T]: ... class accumulate(Iterator[_T], Generic[_T]): if sys.version_info >= (3, 8): @@ -78,12 +81,16 @@ class filterfalse(Iterator[_T], Generic[_T]): 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]: ... + +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: ... diff --git a/tests/stubtest_whitelists/py3_common.txt b/tests/stubtest_whitelists/py3_common.txt index 94ce2ffc0d24..d029986345ff 100644 --- a/tests/stubtest_whitelists/py3_common.txt +++ b/tests/stubtest_whitelists/py3_common.txt @@ -199,12 +199,9 @@ ipaddress._BaseAddress.is_unspecified ipaddress._BaseAddress.max_prefixlen ipaddress._BaseAddress.packed ipaddress._BaseNetwork.max_prefixlen -itertools.accumulate # not a function at runtime -itertools.combinations # not a function at runtime -itertools.groupby # not a function at runtime -itertools.islice # not a function at runtime -itertools.product # not a function at runtime -itertools.repeat # not a function at runtime +itertools.combinations # not a function at runtime, cannot represent with __init__ or __new__ +itertools.groupby # 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__ From 16074baf61efe01ffa555025665fc14ef614dc2d Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Mon, 12 Apr 2021 07:51:31 +0300 Subject: [PATCH 4/5] fix groupby --- stdlib/itertools.pyi | 16 ++++++++++------ tests/stubtest_whitelists/py3_common.txt | 1 - 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index e6f223e725a4..d2ac8b78b3a5 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -77,10 +77,16 @@ class filterfalse(Iterator[_T], Generic[_T]): def __iter__(self) -> Iterator[_T]: ... def __next__(self) -> _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]]]: ... +_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 @@ -109,8 +115,6 @@ class zip_longest(Iterator[Any]): def __iter__(self) -> Iterator[Any]: ... def __next__(self) -> Any: ... -_T1 = TypeVar("_T1") -_T2 = TypeVar("_T2") _T3 = TypeVar("_T3") _T4 = TypeVar("_T4") _T5 = TypeVar("_T5") diff --git a/tests/stubtest_whitelists/py3_common.txt b/tests/stubtest_whitelists/py3_common.txt index d029986345ff..4a3076bfeea8 100644 --- a/tests/stubtest_whitelists/py3_common.txt +++ b/tests/stubtest_whitelists/py3_common.txt @@ -200,7 +200,6 @@ ipaddress._BaseAddress.max_prefixlen ipaddress._BaseAddress.packed ipaddress._BaseNetwork.max_prefixlen itertools.combinations # not a function at runtime, cannot represent with __init__ or __new__ -itertools.groupby # 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 From fc77543dc79ea17788124fd1a703bf8f4118c143 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Mon, 12 Apr 2021 08:36:56 +0300 Subject: [PATCH 5/5] Adding empty overload to count --- stdlib/itertools.pyi | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index d2ac8b78b3a5..b65bf5bda1fa 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -26,7 +26,10 @@ Predicate = Callable[[_T], object] # 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: ... + @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]: ...