Lazy monad

lazy.py

from collections.abc import Callable
from dataclasses import dataclass
from typing import Concatenate, Final


@dataclass(frozen=True)
class Partial1[First, **P, Result]:
    function: Callable[Concatenate[First, P], Result]
    first: First

    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Result:
        return self.function(self.first, *args, **kwargs)


@dataclass(frozen=True)
class Partial2[First, Second, **P, Result]:
    function: Callable[Concatenate[First, Second, P], Result]
    first: First
    second: Second

    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Result:
        return self.function(self.first, self.second, *args, **kwargs)


def identity[Value](value: Value) -> Value:
    return value


def compose[From, Intermediate, To](second: Callable[[Intermediate], To], first: Callable[[From], Intermediate], value: From) -> To:
    return second(first(value))


def make_compose[From, Intermediate, To](second: Callable[[Intermediate], To], first: Callable[[From], Intermediate]) -> Callable[[From], To]:
    return Partial2(compose, second, first)


class Lazy[Result]:
    def __init__[**P](self, function: Callable[P, Result], *args: P.args, **kwargs: P.kwargs) -> None:
        self.function: Final = function
        self.args: Final = args
        self.kwargs: Final = kwargs


def run_lazy[Result](lazy: Lazy[Result]) -> Result:
    return lazy.function(*lazy.args, **lazy.kwargs)


def lazy_identity[Value](value: Value) -> Lazy[Value]:
    return Lazy(lambda: value)


def lazy_map[From, To](function: Callable[[From], To], lazy: Lazy[From]) -> Lazy[To]:
    a: Callable[[Lazy[From]], To] = make_compose(function, run_lazy)
    return Lazy(a, lazy)


def make_lazy_map[From, To](function: Callable[[From], To]) -> Callable[[Lazy[From]], Lazy[To]]:
    return Partial1(lazy_map, function)


def lazy_join[Value](lazy_lazy: Lazy[Lazy[Value]]) -> Lazy[Value]:
    a: Callable[[Lazy[Lazy[Value]]], Value] = make_compose(run_lazy, run_lazy)
    return Lazy(a, lazy_lazy)


def make_lazy_bind[From, To](function: Callable[[From], Lazy[To]]) -> Callable[[Lazy[From]], Lazy[To]]:
    return make_compose(lazy_join, make_lazy_map(function))


def lazy_bind[From, To](function: Callable[[From], Lazy[To]], lazy: Lazy[From]) -> Lazy[To]:
    return make_lazy_bind(function)(lazy)


def make_lazy_compose[From, Intermediate, To](second: Callable[[Intermediate], Lazy[To]], first: Callable[[From], Lazy[Intermediate]]) -> Callable[[From], Lazy[To]]:
    return make_compose(make_lazy_bind(second), first)


def lazy_compose[From, Intermediate, To](second: Callable[[Intermediate], Lazy[To]], first: Callable[[From], Lazy[Intermediate]], value: From) -> Lazy[To]:
    return make_lazy_compose(second, first)(value)


type Stream[Value] = tuple[Value, Lazy[Stream[Value]]]


def take[Value](number: int, stream: Stream[Value]) -> list[Value]:
    return [stream[0], *take(number - 1, run_lazy(stream[1]))] if number else []


def stream_of[Value](value: Value) -> Stream[Value]:
    return value, Lazy(stream_of, value)


stream_of_ones = stream_of(1)


print(take(10, stream_of_ones))