arrows.py

# An implementation of Haskell's arrow library in Python, done as an exercise

from __future__ import annotations
from typing import Callable

class Arrow:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args):
        return self.func(*args)

    def and_then(self, other: Arrow) -> Arrow:
        "Implements >>>"
        # a b c -> a c d -> a b d
        return Arrow(lambda x: other(self(x)))

    def first(self) -> Arrow:
        return Arrow(lambda b, d: (self(b), d))

    def second(self) -> Arrow:
        return Arrow(lambda d, b: (d, self(b)))

    def times(self, other: Arrow) -> Arrow:
        "Implements ***"
        return self.first().and_then(other.second())

    def join(self, other: Arrow) -> Arrow:
        "Implements &&&"
        return split().and_then(self.times(other))

    def lift_a2(self, other: Arrow, binop: Callable) -> Arrow:
        return split().and_then(self.first().and_then(other.second().and_then(unsplit(binop))))

def split() -> Arrow:
    return Arrow(lambda x: (x, x))

def unsplit(func: Callable) -> Arrow:
    return Arrow(func)