sm.py

"""
Example state machine with conditions.
"""

from functools import partialmethod


" Possible states."
WORK = 'Work'
TRAVEL = 'Travel'
HOME = 'Home'

" Conditions (actions)."
LEAVE = 'Leave current place'
DO_WORK = 'Do work'
SLEEP = 'Sleep'
GO_HOME = 'Go to home'
GO_WORK = 'Go to Work'

" State transitions."
TRANSITIONS = {
    WORK: {
        DO_WORK: WORK,
        LEAVE: TRAVEL,
    },
    TRAVEL: {
        GO_HOME: HOME,
        GO_WORK: WORK,
    },
    HOME: {
        SLEEP: HOME,
        LEAVE: TRAVEL,
    },
}


class Context:
    """
    The wrapper object.
    """

    def __init__(self):
        self.state = HOME

    def transition(self, condition):
        """ Try to transition to a new state."""
        try:
            self._transition_state(condition)
        except KeyError:
            raise ValueError('Invalid condition')

    def _transition_state(self, condition):
        """ Transition to a new state."""
        self.state = TRANSITIONS[self.state][condition]


if __name__ == "__main__":
    context = Context()
    assert context.state == HOME
    context.transition(LEAVE)
    assert context.state == TRAVEL
    try:
        context.transition(SLEEP)
        assert False, "Can't sleep on the street"
    except ValueError:
        assert True
    context.transition(GO_HOME)
    assert context.state == HOME