comp.py

from functools import partial
from typing import Callable, TypeVar, overload, List, Sequence

S = TypeVar('S')
T = TypeVar('T')
U = TypeVar('U')


def comp(g, f):
    # type: (Callable[[T], U], Callable[[S], T]) -> Callable[[S], U]
    def composed(x):
        # type : (S) -> U
        return g(f(x))
    return composed


@overload
def cfilter(f):
    # type: (Callable[[T], bool]) -> Callable[[Sequence[T]], Sequence[T]]
    pass


@overload
def cfilter(f, lst):
    # type: (Callable[[T], bool], Sequence[T]) -> Sequence[T]
    pass


def cfilter(f, lst=None):
    if lst is None:
        return partial(cfilter, f)
    return filter(f, lst)


@overload
def cmap(f):
    # type: (Callable[[T], S]) -> Callable[[Sequence[T]], Sequence[S]]
    pass


@overload
def cmap(f, lst):
    # type: (Callable[[T], S], Sequence[T]) -> Sequence[S]
    pass


def cmap(f, lst=None):
    if lst is None:
        return partial(cmap, f)
    return map(f, lst)


def test():
    # type: () -> None
    x = ["a", "b", "c"]  # type: Sequence[str]

    mapfilter = comp(cfilter(lambda x: x < 10),
                     cmap(len))

    y = mapfilter(x)
    reveal_type(y)

    def lt10(x):
        # type: (int) -> bool
        return x < 10

    mapfilter2 = comp(cfilter(lt10), cmap(lambda s: len(s) * 2))
    reveal_type(mapfilter2)
    y2 = mapfilter2(x)
    reveal_type(y2)

    comp(lambda x: x + 1, lambda s: len(s))
    # Sadly, one can't do this:
    # comp(
    #     lambda x: x + 1,  # type: (int) -> int
    #     len
    # )