group_by.py

from collections import defaultdict
from typing import Callable, Dict, Iterable, List, Optional

T = TypeVar("T")
U = TypeVar("U")


def group_by(
    func: Callable[[T], U],
    coll: Iterable[T],
    xform: Optional[Callable[[T], Any]] = None,
) -> Dict[U, List]:
    """
    Returns a dictionary of keys, generated via the supplied func, to lists of items
    from the supplied collection.
    
    If desired, these items can be transformed before they are added to the dictionary.
    """
    ddict: Dict[U, List] = defaultdict(list)
    for item in coll:
        ddict[func(item)].append(xform(item) if xform is not None else item)
    return ddict


class TestGroupBy:
    def test_without_xform(self):
        assert {"> 5": [7, 9, 6], "< 5": [4, 1, 2]} == group_by(
            lambda i: "> 5" if i > 5 else "< 5", [4, 7, 1, 2, 9, 6]
        )

    def test_with_xform(self):
        assert {"> 5": ["7", "9", "6"], "< 5": ["4", "1", "2"]} == group_by(
            lambda i: "> 5" if i > 5 else "< 5", [4, 7, 1, 2, 9, 6], str
        )