linq.py

from typing import (Callable, Dict, Generic, Iterable, List, Protocol, Tuple,
                    TypeVar, Union)

T = TypeVar("T")
S = TypeVar("S")
U = TypeVar("U", contravariant=True)


class SupportsDunderLT(Protocol[U]):
    def __lt__(self, __other: U) -> bool:
        ...


class SupportsDunderGT(Protocol[U]):
    def __gt__(self, __other: U) -> bool:
        ...


class Linq(Generic[T]):
    def __init__(self, iterable: Iterable[T]) -> None:
        self._iterable = iterable

    def select(self, func: Callable[[T], S]) -> "Linq[S]":
        return Linq(map(func, self._iterable))

    def where(self, func: Callable[[T], bool]) -> "Linq[T]":
        return Linq(filter(func, self._iterable))

    def group_by(self, func: Callable[[T], S]) -> "Linq[Tuple[S, Linq[T]]]":
        groups: Dict[S, List[T]] = {}
        for item in self._iterable:
            key = func(item)
            if key not in groups:
                groups[key] = []
            groups[key].append(item)
        return Linq((key, Linq(group)) for key, group in groups.items())

    def order_by(
        self, func: Callable[[T], Union[SupportsDunderLT, SupportsDunderGT]]
    ) -> "Linq[T]":
        return Linq(sorted(self._iterable, key=func))

    def to_list(self) -> List[T]:
        return list(self._iterable)

    def __iter__(self) -> Iterable[T]:
        return iter(self._iterable)


if __name__ == "__main__":
    linq = Linq([3, 4, 5, 9, 2, 6, 1, 7, 8])
    print(linq.select(lambda x: x * 2).where(lambda x: x > 5).to_list())
    print(linq.group_by(lambda x: x % 2).select(lambda x: x[1].to_list()).to_list())
    print(linq.order_by(lambda x: x).to_list())