23 lines
591 B
Python
23 lines
591 B
Python
import itertools
|
|
from collections.abc import Callable, Iterable, Iterator
|
|
from typing import Protocol, Self, TypeVar
|
|
|
|
|
|
class SupportsLessThan(Protocol):
|
|
def __lt__(self, other: Self) -> bool:
|
|
raise NotImplementedError()
|
|
|
|
|
|
T = TypeVar("T")
|
|
KeyT = TypeVar("KeyT", bound=SupportsLessThan)
|
|
GroupT = TypeVar("GroupT")
|
|
|
|
|
|
def groupby(
|
|
it: Iterable[T],
|
|
key_fn: Callable[[T], KeyT],
|
|
group_fn: Callable[[Iterable[T]], GroupT] = lambda x: x,
|
|
) -> Iterator[tuple[KeyT, GroupT]]:
|
|
for k, g in itertools.groupby(sorted(it, key=key_fn), key=key_fn):
|
|
yield k, group_fn(g)
|