Python functional operator chaining

chain.py

from __future__ import annotations

import abc
import typing as t


class Operator[A, B](abc.ABC):
    __slots__ = ()

    @abc.abstractmethod
    def __call__(self, item: A) -> t.Iterator[B]: ...

    @t.final
    def compose[C](self, other: Operator[B, C]) -> Operator[A, C]:
        return _OperatorPair[A, C](self, other)

    @t.final
    def apply(self, data: t.Iterable[A]) -> Operation[A, B]:
        return Operation(data, self)

    __matmul__ = t.final(compose)
    __rmatmul__ = t.final(apply)


@t.final
class FunctionalOperator[A, B](Operator[A, B]):
    __slots__ = ("func",)

    def __init__(self, func: t.Callable[[A], t.Iterator[B]]) -> None:
        self.func = func

    def __call__(self, item: A) -> t.Iterator[B]:
        return self.func(item)


@t.final
class _OperatorPair[A, B, C](Operator[A, C]):
    __slots__ = ("left", "right")

    def __init__(self, left: Operator[A, B], right: Operator[B, C]) -> None:
        self.left = left
        self.right = right

    def __call__(self, item: A) -> t.Iterator[C]:
        for b in self.left(a):
            yield from self.right(b)


class Operation[A, B]:
    __slots__ = ("data", "operator")

    def __init__(self, data: t.Iterator[A], operator: Operator[A, B]) -> None:
        self.data = data
        self.operator = operator

    def __iter__(self) -> t.Iterator[B]:
        for item in self.data:
            yield from self.operator(item)

    def compose[C](self, other: Operator[B, C]) -> Operation[A, C]:
        return Operation[A, C](self.data, self.operator.compose(other))


def ch_flat_map[A, B](func: t.Callable[[A], t.Iterator[B]]) -> Operator[A, B]:
    return FunctionalOperator(func)


def ch_map[A, B](func: t.Callable[[A], t.Iterator[B]]) -> Operator[A, B]:
    @ch_flat_map
    def ch_map_impl(item: A) -> t.Iterator[B]:
        yield func(item)

    return ch_map_impl


def ch_filter[A, B](func: t.Callable[[A], bool]) -> Operator[A, A]:
    @ch_flat_map
    def ch_filter_impl(item: A) -> t.Iterator[A]:
        if func(item):
            yield item

    return ch_filter_impl


def fizz_buzz(i: int) -> t.Iterator[str]:
    yield str(i)
    if i % 3 == 0:
        yield "fizz"
    if i % 5 == 0:
        yield "buzz"


items = range(100) \
    @ ch_map(lambda i: i + 1) \
    @ ch_filter(lambda i: i % 3 == 0 or i % 5 == 0) \
    @ ch_flat_map(fizz_buzz)


for item in items:
    print(item, end=" ")