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

1# mypy: disable_error_code="operator" 

2from __future__ import annotations 

3 

4from collections.abc import Callable 

5 

6 

7class RWS[R, X, S, T]: 

8 """Reader + Writer + State monad — environment, output, and mutable state. 

9 

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

11 

12 Type parameters: 

13 R: Environment type. 

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

15 S: State type. 

16 T: Result type. 

17 """ 

18 

19 __slots__ = ('_run',) 

20 

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

22 self._run = run 

23 

24 def map[U](self, f: Callable[[T], U]) -> RWS[R, X, S, U]: 

25 """Apply a pure function to the result.""" 

26 def run(r: R, s: S) -> tuple[U, S, X]: 

27 a, s, x = self._run(r, s) 

28 return f(a), s, x 

29 

30 return RWS(run) 

31 

32 def and_then[U](self, f: Callable[[T], RWS[R, X, S, U]]) -> RWS[R, X, S, U]: 

33 """Monadic bind — chain, threading state and concatenating outputs.""" 

34 def run(r: R, s: S) -> tuple[U, S, X]: 

35 a, s1, x1 = self._run(r, s) 

36 b, s2, x2 = f(a)(r, s1) 

37 return b, s2, x1 + x2 

38 

39 return RWS(run) 

40 

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

42 """Execute the computation with environment and initial state.""" 

43 return self._run(r, state) 

44 

45 @staticmethod 

46 def tell(x: X) -> RWS[R, X, S, None]: 

47 """Append to the writer output.""" 

48 return RWS(lambda r, s: (None, s, x)) 

49 

50 @staticmethod 

51 def pure(value: T, monoid: X) -> RWS[R, X, S, T]: 

52 """Wrap a pure value with initial writer output.""" 

53 return RWS(lambda r, s: (value, s, monoid)) 

54 

55 @staticmethod 

56 def get(monoid: X) -> RWS[R, X, S, T]: 

57 """Return the current state as value.""" 

58 return RWS(lambda r, s: (s, s, monoid)) 

59 

60 @staticmethod 

61 def put(s: S, monoid: X) -> RWS[R, X, S, None]: 

62 """Replace the current state.""" 

63 return RWS(lambda r, _: (None, s, monoid))