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

1from __future__ import annotations 

2 

3from collections.abc import Callable 

4 

5 

6class RS[R, S, T]: 

7 """Reader + State monad — computation with environment and mutable state. 

8 

9 Wraps a function ``(R, S) -> (T, S)``. Threads state through the chain. 

10 

11 Type parameters: 

12 R: Environment type. 

13 S: State type. 

14 T: Result type. 

15 

16 Examples: 

17 :: 

18 

19 rs = RS(lambda r, s: (r + s, s + 1)) 

20 assert rs(10, 5) == (15, 6) 

21 

22 chained = rs.and_then(lambda v: RS(lambda r, s: (v * 10, s + 1))) 

23 assert chained(2, 3) == (50, 5) # (2+3=5, s=4) → (5*10=50, s=5) 

24 """ 

25 

26 __slots__ = ('_run',) 

27 

28 def __init__(self, function: Callable[[R, S], tuple[T, S]]) -> None: 

29 self._run = function 

30 

31 def map[U](self, f: Callable[[T], U]) -> RS[R, S, U]: 

32 """Apply a pure function to the result, preserving state.""" 

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

34 a, s = self._run(r, s) 

35 return f(a), s 

36 

37 return RS(run) 

38 

39 def and_then[U](self, f: Callable[[T], RS[R, S, U]]) -> RS[R, S, U]: 

40 """Monadic bind — chain, threading state through both steps.""" 

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

42 a, s1 = self._run(r, s) 

43 b, s2 = f(a)(r, s1) 

44 return b, s2 

45 

46 return RS(run) 

47 

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

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

50 return self._run(r, state) 

51 

52 @staticmethod 

53 def pure(value: T) -> RS[R, S, T]: 

54 """Wrap a pure value, passing state through unchanged.""" 

55 return RS(lambda r, s: (value, s)) 

56 

57 @staticmethod 

58 def get() -> RS[R, S, T]: 

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

60 return RS(lambda r, s: (s, s)) 

61 

62 @staticmethod 

63 def put(s: S) -> RS[R, S, None]: 

64 """Replace the current state.""" 

65 return RS(lambda r, _: (None, s))