diff --git a/stdlib/itertools.pyi b/stdlib/itertools.pyi index 55f9c92d8..881fb236b 100644 --- a/stdlib/itertools.pyi +++ b/stdlib/itertools.pyi @@ -10,6 +10,7 @@ _T = TypeVar("_T") _S = TypeVar("_S") _N = TypeVar("_N", int, float, SupportsFloat, SupportsInt, SupportsIndex, SupportsComplex) _T_co = TypeVar("_T_co", covariant=True) +_S_co = TypeVar("_S_co", covariant=True) _T1 = TypeVar("_T1") _T2 = TypeVar("_T2") _T3 = TypeVar("_T3") @@ -84,13 +85,13 @@ class filterfalse(Iterator[_T]): def __iter__(self) -> Self: ... def __next__(self) -> _T: ... -class groupby(Iterator[tuple[_T, Iterator[_S]]], Generic[_T, _S]): +class groupby(Iterator[tuple[_T_co, Iterator[_S_co]]], Generic[_T_co, _S_co]): @overload def __new__(cls, iterable: Iterable[_T1], key: None = None) -> groupby[_T1, _T1]: ... @overload def __new__(cls, iterable: Iterable[_T1], key: Callable[[_T1], _T2]) -> groupby[_T2, _T1]: ... def __iter__(self) -> Self: ... - def __next__(self) -> tuple[_T, Iterator[_S]]: ... + def __next__(self) -> tuple[_T_co, Iterator[_S_co]]: ... class islice(Iterator[_T]): @overload diff --git a/test_cases/stdlib/itertools/check_itertools_recipes.py b/test_cases/stdlib/itertools/check_itertools_recipes.py index b74d92ea5..340811eec 100644 --- a/test_cases/stdlib/itertools/check_itertools_recipes.py +++ b/test_cases/stdlib/itertools/check_itertools_recipes.py @@ -9,7 +9,7 @@ import collections import math import operator import sys -from itertools import chain, combinations, count, cycle, filterfalse, islice, repeat, starmap, tee, zip_longest +from itertools import chain, combinations, count, cycle, filterfalse, groupby, islice, repeat, starmap, tee, zip_longest from typing import Any, Callable, Hashable, Iterable, Iterator, Sequence, Tuple, Type, TypeVar, Union, overload from typing_extensions import Literal, TypeAlias, TypeVarTuple, Unpack @@ -272,6 +272,15 @@ def unique_everseen(iterable: Iterable[_T], key: Callable[[_T], Hashable] | None # yield from dict(zip(map(key, t1), t2)).values() +# Slightly adapted from the docs recipe; a one-liner was a bit much for pyright +def unique_justseen(iterable: Iterable[_T], key: Callable[[_T], bool] | None = None) -> Iterator[_T]: + "List unique elements, preserving order. Remember only the element just seen." + # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B + # unique_justseen('ABBcCAD', str.lower) --> A B c A D + g: groupby[_T | bool, _T] = groupby(iterable, key) + return map(next, map(operator.itemgetter(1), g)) + + def powerset(iterable: Iterable[_T]) -> Iterator[tuple[_T, ...]]: "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)" s = list(iterable)