<NME> state.py
<BEF> from typing import Callable, Tuple, Any, TypeVar, Generic

from .util import Unit
from .typing import Functor
from .typing import Monad

TState = TypeVar("TState")
TSource = TypeVar("TSource")
TResult = TypeVar("TResult")


class State(Generic[TSource, TState]):
    """The state monad.

    Wraps stateful computations. A stateful computation is a function
    that takes a state and returns a result and new state:

        state -> (result, state')
    """

    def __init__(self, fn: Callable[[TState], Tuple[TSource, TState]]) -> None:
        """Initialize a new state.

        Keyword arguments:
        fn -- State processor.
        """

        self._fn = fn

    @classmethod
    def unit(cls, value: TSource) -> "State[TSource, TState]":
        r"""Create new State.

        The unit function creates a new State object wrapping a stateful
        computation.

        State $ \s -> (x, s)
        """
        def _(a, state):
            return mapper(a), state

        return State(lambda state: _(*self.run()(state)))

    def bind(self, fn: Callable[[Any], 'State']) -> 'State':
        r"""m >>= k = State $ \s -> let (a, s') = runState m s

        """

        def _(result, state):
            return fn(result).run()(state)

        return State(lambda state: _(*self.run()(state)))

    @classmethod
    def get(cls) -> 'State':

    @classmethod
    def get(cls) -> "State[TState, TState]":
        r"""get = state $ \s -> (s, s)"""
        return State(lambda state: (state, state))
        r"""put newState = state $ \s -> ((), newState)"""
        return State(lambda state: (Unit, new_state))

    def run(self):
        return self._get_value()

    def __call__(self, state):
        return self.run()(state)

    def __eq__(self, other) -> bool:
        """Test if two stateful computations are equal.
assert issubclass(State, Functor)
assert issubclass(State, Monad)

<MSG> Accept args for run() to avoid run()(*args)

<DFF> @@ -39,7 +39,7 @@ class State(Monad, Functor):
         def _(a, state):
             return mapper(a), state
 
-        return State(lambda state: _(*self.run()(state)))
+        return State(lambda state: _(*self.run(state)))
 
     def bind(self, fn: Callable[[Any], 'State']) -> 'State':
         r"""m >>= k = State $ \s -> let (a, s') = runState m s
@@ -47,9 +47,9 @@ class State(Monad, Functor):
         """
 
         def _(result, state):
-            return fn(result).run()(state)
+            return fn(result).run(state)
 
-        return State(lambda state: _(*self.run()(state)))
+        return State(lambda state: _(*self.run(state)))
 
     @classmethod
     def get(cls) -> 'State':
@@ -61,11 +61,17 @@ class State(Monad, Functor):
         r"""put newState = state $ \s -> ((), newState)"""
         return State(lambda state: (Unit, new_state))
 
-    def run(self):
-        return self._get_value()
+    def run(self, *args):
+        """Return wrapped state computation.
+
+        This is the inverse of unit and returns the wrapped function.
+        If we receive args, we call the function directly to avoid the
+        ugly `run()(args)` pattern.
+        """
+        return self._get_value()(*args) if args else self._get_value()
 
     def __call__(self, state):
-        return self.run()(state)
+        return self.run(state)
 
     def __eq__(self, other) -> bool:
         """Test if two stateful computations are equal.

Accept args for run() to avoid run()(*args)