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

from .util import Unit
from .abc import Functor
from .abc import Monad

A = TypeVar('A')
B = TypeVar('B')
S = TypeVar('S')


class State(Generic[A, S], Monad['State[S]'], Functor['State[S]']):

    """The state monad.

    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)
        """
        return cls(lambda state: (value, state))

    def map(self, mapper: Callable[[TSource], TResult]) -> "State[TResult, TState]":
        def _(a: Any, state: Any) -> Tuple[Any, Any]:
            return mapper(a), state

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

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

        def _(result: Any, state: Any) -> Tuple[Any, Any]:
            return fn(result).run(state)

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

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

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

    def run(self, state: TState) -> Tuple[TSource, TState]:
        """Return wrapped state computation.

        This is the inverse of unit and returns the wrapped function.
        """
        return self._fn(state)

    def __call__(self, state: Any) -> Tuple:
        return self.run(state)


assert issubclass(State, Functor)
assert issubclass(State, Monad)

<MSG> Fixes for type hints

<DFF> @@ -4,12 +4,12 @@ from .util import Unit
 from .abc import Functor
 from .abc import Monad
 
-A = TypeVar('A')
-B = TypeVar('B')
-S = TypeVar('S')
+A = TypeVar('A')  # Type of value
+B = TypeVar('B')  # Type of new value
+S = TypeVar('S')  # Type of state
 
 
-class State(Generic[A, S], Monad['State[S]'], Functor['State[S]']):
+class State(Generic[A, S], Monad[A], Functor[A]):
 
     """The state monad.
 

Fixes for type hints