During a pair programming we found out it is pretty easy to implement the | operator in Python. In the previous post, I implemented a function that can transform an input value using a sequence of functions. Using operator overriding, we can do the same using a fancy pipe operator syntax.

With the current implementation I have, I can do something like this.

add_1: Fn[int, int] = lambda x: x + 1
multiply_by: Fn[int, Fn[int, int]] = lambda y: lambda x: x * y

# + assume all the replicate, fmap and flatten exist in the module A

result = (
    start_pipe(1)
    | A.replicate(3)
    | A.fmap(add_1)
    | A.replicate(2)
    | A.flatten
    | A.fmap(multiply_by(2))
    | end_pipe
)
assert result == [4, 4, 4, 4, 4, 4]

The function start_pipe converts a value of type A to a wrapper object Pipeable[A]. The trick is to override __or__ method on the object and expect a fn: Fn[A, B] as an input argument. Every invocation of the pipeable | fn expression than creates a new Pipeable object.

The result of p(a) | fn1 | fn2 | fn3 ... is always Pipeable. To actually evaluate the pipeable object, I introduced a secret class _EndPipe and a single instance end_pipe: _EndPipe. In the implementation of the __or__ method I only need to check whether the instance of _EndPipe has been provided and if so the pipeable is evaluated instead and the result is returned. I created overloads for the __or__ method to make it type-check and that's it!

from __future__ import annotations

from functools import reduce
from typing import Any, Generic, TypeGuard, TypeVar, cast, overload

from dogs.function import Fn

A = TypeVar("A")
B = TypeVar("B")

class _EndPipe:
    pass

end_pipe = _EndPipe()

class Pipeable(Generic[A]):
    def __init__(self, a: A) -> None:
        self._a = a
        self._fns: list[Fn[Any, Any]] = []

    def __with_new_fn(self, fn: Fn[A, B]) -> Pipeable[B]:
        new_pipeable = Pipeable(self._a)
        new_pipeable._fns = self._fns
        new_pipeable._fns.append(fn)
        return cast(Pipeable[B], new_pipeable)

    @staticmethod
    def __is_end_pipe(symbol: Any) -> TypeGuard[_EndPipe]:
        return isinstance(symbol, _EndPipe)

    @overload
    def __or__(self, fn: Fn[A, B]) -> Pipeable[B]:
        ...

    @overload
    def __or__(self, fn: _EndPipe) -> A:
        ...

    def __or__(self, fn: Fn[A, B] | _EndPipe) -> Pipeable[B] | A:
        if self.__is_end_pipe(fn):
            return self.eval()

        return self.__with_new_fn(cast(Fn[A, B], fn))

    def eval(self) -> A:
        return reduce(lambda acc, fn: fn(acc), self._fns, self._a)


def start_pipe(a: A) -> Pipeable[A]:
    return Pipeable(a)