Properly specify groupby

Fixes #35

functionalpy/_sequence.py

@@ -1,63 +1,71 @@
 import itertools
-from collections.abc import Callable, Iterable, Iterator, Sequence
+from collections import defaultdict
+from collections.abc import (
+    Callable,
+    Iterable,
+    Iterator,
+    Mapping,
+    Sequence,
+)
 from dataclasses import dataclass
 from functools import reduce
 from typing import Generic, TypeVar
 
-_T0 = TypeVar("_T0")
-_T1 = TypeVar("_T1")
+_T = TypeVar("_T0")
+_S = TypeVar("_T1")
 
 
 @dataclass(frozen=True)
-class Group(Generic[_T0]):
+class Group(Generic[_T]):
     key: str
-    value: "Seq[_T0]"
+    value: "Seq[_T]"
 
 
-class Seq(Generic[_T0]):
-    def __init__(self, iterable: Iterable[_T0]):
+class Seq(Generic[_T]):
+    def __init__(self, iterable: Iterable[_T]):
         self._seq = iterable
 
     ### Reductions
     def count(self) -> int:
         return sum(1 for _ in self._seq)
 
     ### Output
-    def to_list(self) -> list[_T0]:
+    def to_list(self) -> list[_T]:
         return list(self._seq)
 
-    def to_tuple(self) -> tuple[_T0, ...]:
+    def to_tuple(self) -> tuple[_T, ...]:
         return tuple(self._seq)
 
-    def to_iter(self) -> Iterator[_T0]:
+    def to_iter(self) -> Iterator[_T]:
         return iter(self._seq)
 
-    def to_set(self) -> set[_T0]:
+    def to_set(self) -> set[_T]:
         return set(self._seq)
 
     ### Transformations
     def map(  # noqa: A003 # Ignore that it's shadowing a python built-in
         self,
-        func: Callable[[_T0], _T1],
-    ) -> "Seq[_T1]":
+        func: Callable[[_T], _S],
+    ) -> "Seq[_S]":
         return Seq(map(func, self._seq))
 
-    def filter(self, func: Callable[[_T0], bool]) -> "Seq[_T0]":  # noqa: A003
+    def filter(self, func: Callable[[_T], bool]) -> "Seq[_T]":  # noqa: A003
         return Seq(filter(func, self._seq))
 
-    def reduce(self, func: Callable[[_T0, _T0], _T0]) -> _T0:
+    def reduce(self, func: Callable[[_T, _T], _T]) -> _T:
         return reduce(func, self._seq)
 
-    def group_by(
-        self, func: Callable[[_T0], str]
-    ) -> "Seq[Group[_T0]]":
-        result = (
-            Group(key=key, value=Seq(value))
-            for key, value in itertools.groupby(self._seq, key=func)
-        )
-        return Seq(result)
+    def groupby(
+        self, func: Callable[[_T], str]
+    ) -> "Mapping[str, Sequence[_T]]":
+        mapping: defaultdict[str, list[_T]] = defaultdict(list)
+
+        for item in self._seq:
+            mapping[func(item)].append(item)
+
+        return dict(mapping)
 
-    def flatten(self) -> "Seq[_T0]":
+    def flatten(self) -> "Seq[_T]":
         return Seq(
             item
             for sublist in self._seq

functionalpy/_sequence.pyi

@@ -1,36 +1,42 @@
 """
 This type stub file was generated by pyright.
+# TODO: https://github.com/MartinBernstorff/FunctionalPy/issues/38 ci: ensure .pyi files are aligned with implementation
 """
 
-from collections.abc import Callable, Iterable, Iterator
+from collections.abc import (
+    Callable,
+    Iterable,
+    Iterator,
+    Mapping,
+    Sequence,
+)
 from dataclasses import dataclass
 from typing import Generic, TypeVar, overload
 
-_T0 = TypeVar("_T0")
-_T1 = TypeVar("_T1")
+_T = TypeVar("_T")
 _S = TypeVar("_S")
 
 @dataclass(frozen=True)
-class Group(Generic[_T0]):
+class Group(Generic[_T]):
     key: str
-    value: Seq[_T0]
+    value: Seq[_T]
     ...
 
-class Seq(Generic[_T0]):
-    def __init__(self, iterable: Iterable[_T0]) -> None: ...
+class Seq(Generic[_T]):
+    def __init__(self, iterable: Iterable[_T]) -> None: ...
     def count(self) -> int: ...
-    def to_list(self) -> list[_T0]: ...
-    def to_tuple(self) -> tuple[_T0, ...]: ...
-    def to_iter(self) -> Iterator[_T0]: ...
-    def to_set(self) -> set[_T0]: ...
-    def map(self, func: Callable[[_T0], _T1]) -> Seq[_T1]:  # noqa: A003
+    def to_list(self) -> list[_T]: ...
+    def to_tuple(self) -> tuple[_T, ...]: ...
+    def to_iter(self) -> Iterator[_T]: ...
+    def to_set(self) -> set[_T]: ...
+    def map(self, func: Callable[[_T], _S]) -> Seq[_S]:  # noqa: A003
         ...
-    def filter(self, func: Callable[[_T0], bool]) -> Seq[_T0]:  # noqa: A003
+    def filter(self, func: Callable[[_T], bool]) -> Seq[_T]:  # noqa: A003
         ...
-    def reduce(self, func: Callable[[_T0, _T0], _T0]) -> _T0: ...
-    def group_by(
-        self, func: Callable[[_T0], str]
-    ) -> Seq[Group[_T0]]: ...
+    def reduce(self, func: Callable[[_T, _T], _T]) -> _T: ...
+    def groupby(
+        self, func: Callable[[_T], str]
+    ) -> Mapping[str, Sequence[_T]]: ...
     @overload
     def flatten(self: Seq[list[_S]]) -> Seq[_S]: ...
     @overload

functionalpy/test_sequence.py

@@ -29,6 +29,28 @@ def test_reduce():
     assert result == 3
 
 
+class TestGroupby:
+    def test_grouped_filter(self):
+        sequence = Seq([1, 2, 3, 4])
+
+        def is_even(num: int) -> str:
+            if num % 2 == 0:
+                return "even"
+            return "odd"
+
+        grouped = sequence.groupby(is_even)
+        assert grouped == {
+            "odd": [1, 3],
+            "even": [2, 4],
+        }
+
+    def test_grouped_map(self):
+        ...
+
+    def test_grouped_reduce(self):
+        ...
+
+
 def test_flatten():
     test_input = ((1, 2), (3, 4))
     sequence = Seq(test_input)