57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""A Writer Monad. Requires Python >= 3.7 for type hints."""
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import math
|
|
import os
|
|
|
|
from typing import Any
|
|
from typing import Callable
|
|
from typing import Generic
|
|
from typing import List
|
|
from typing import TypeVar
|
|
from typing import Union
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class Writer(Generic[T]):
|
|
def __init__(self, value: Union[T, Writer[T]], *msgs: str):
|
|
if isinstance(value, Writer):
|
|
self.value: T = value.value
|
|
self.msgs: List[str] = value.msgs + list(msgs)
|
|
else:
|
|
self.value = value
|
|
self.msgs = list(f"{msg}: {self.value}" for msg in msgs)
|
|
|
|
def bind(self, func: Callable[[T], Writer[Any]]) -> Writer[Any]:
|
|
writer = func(self.value)
|
|
return Writer(writer, *self.msgs)
|
|
|
|
def __rshift__(self, func: Callable[[T], Writer[Any]]) -> Writer[Any]:
|
|
return self.bind(func)
|
|
|
|
def __str__(self):
|
|
return f"{self.value}\n{os.linesep.join(reversed(self.msgs))}"
|
|
|
|
def __repr__(self):
|
|
return f"Writer({self.value}, \"{', '.join(reversed(self.msgs))}\")"
|
|
|
|
|
|
def lift(func: Callable, msg: str) -> Callable[[Any], Writer[Any]]:
|
|
"""Return a writer monad version of the simple function `func`."""
|
|
|
|
@functools.wraps(func)
|
|
def wrapped(value):
|
|
return Writer(func(value), msg)
|
|
|
|
return wrapped
|
|
|
|
|
|
if __name__ == "__main__":
|
|
square_root = lift(math.sqrt, "square root")
|
|
add_one = lift(lambda x: x + 1, "add one")
|
|
half = lift(lambda x: x / 2, "div two")
|
|
|
|
print(Writer(5, "initial") >> square_root >> add_one >> half)
|