Coverage for greyhorse / river / private / monads / rw.py: 100%

1# mypy: disable_error_code="operator" 

2from __future__ import annotations 

3 

4from collections.abc import Callable 

5 

6 

7class ReaderWriter[R, X, T]: 

8 """Reader + Writer monad — computation with environment and accumulated output. 

9 

10 Wraps ``R -> (T, X)``. Writer outputs are combined with ``+`` in ``and_then``. 

11 

12 Type parameters: 

13 R: Environment type. 

14 X: Writer output type (must support ``+``). 

15 T: Result type. 

16 

17 Examples: 

18 :: 

19 

20 rw = ReaderWriter(lambda r: (r * 2, [r])) 

21 assert rw(5) == (10, [5]) 

22 

23 chained = rw.and_then( 

24 lambda v: ReaderWriter(lambda r: (v + 1, [r * 10])) 

25 ) 

26 assert chained(5) == (11, [5, 50]) # outputs concatenated 

27 """ 

28 

29 __slots__ = ('_run',) 

30 

31 def __init__(self, run: Callable[[R], tuple[T, X]]) -> None: 

32 self._run = run 

33 

34 def map[U](self, f: Callable[[T], U]) -> ReaderWriter[R, X, U]: 

35 """Apply a pure function to the result, preserving writer output.""" 

36 def run(r: R) -> tuple[U, X]: 

37 a, x = self._run(r) 

38 return f(a), x 

39 

40 return ReaderWriter(run) 

41 

42 def and_then[U](self, f: Callable[[T], ReaderWriter[R, X, U]]) -> ReaderWriter[R, X, U]: 

43 """Monadic bind — chain, concatenating writer outputs with ``+``.""" 

44 def run(r: R) -> tuple[U, X]: 

45 a, x1 = self._run(r) 

46 b, x2 = f(a)(r) 

47 return b, x1 + x2 

48 

49 return ReaderWriter(run) 

50 

51 def __call__(self, r: R) -> tuple[T, X]: 

52 """Execute the computation with the given environment.""" 

53 return self._run(r)