diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index aef3aac78635..b65bf5bda1fa 100644 --- a/stdlib/itertools.pyi +++ b/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, ...]: ... diff --git a/tests/stubtest_whitelists/py36.txt b/tests/stubtest_whitelists/py36.txt index 7f2e683b2caf..95807fe259ee 100644 --- a/tests/stubtest_whitelists/py36.txt +++ b/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 diff --git a/tests/stubtest_whitelists/py37.txt b/tests/stubtest_whitelists/py37.txt index dcd3585c884e..973768a431c9 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 # Added in 3.8 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 af577308f4e0..4a3076bfeea8 100644 --- a/tests/stubtest_whitelists/py3_common.txt +++ b/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__