FSM

fsm.py

import prettytable
import six
import sys


class FSM(object):

    """A basic state machine."""
    def __init__(self, start_state):
        self._table = {}
        self._start_state = start_state
        self._current = (None, None)

    @property
    def start_state(self):
        return self._start_state

    @property
    def current_state(self):
        return self._current[0]

    def add_state(self, state):
        if state not in self._table:
            self._table[state] = {}

    def add_transition(self, start, end, event, on_enter=None, on_exit=None):
        if start not in self._table:
            raise ValueError("Can not add a transition that starts in a"
                             " undefined state '%s'" % start)
        if end not in self._table:
            raise ValueError("Can not add a transition that ends in a"
                             " undefined state '%s'" % end)
        self._table[start][event] = (end, on_enter, on_exit)

    def process_event(self, event):
        current_state, old_on_exit = self._current
        end, on_enter, new_on_exit = self._table[current_state][event]
        if old_on_exit is not None:
            old_on_exit(event)
        if on_enter is not None:
            on_enter(event)
        self._current_state = (end, new_on_exit)

    def start(self):
        if self._start_state not in self._table:
            raise RuntimeError("Can not start from a undefined state")
        self.reset()

    def reset(self):
        self._current = (self._start_state, None)

    def pformat(self):
        tbl = prettytable.PrettyTable(["State", "Event", "End",
                                       "On Enter", "On Exit"])
        for state in sorted(six.iterkeys(self._table)):
            for event in sorted(six.iterkeys(self._table[state])):
                end, on_enter, on_exit = self._table[state][event]
                tbl.add_row([state, event, end, on_enter, on_exit])
        return tbl.get_string()
    

def print_enter(e):
    print("%s occurred" % (e))


def print_exit(e):
    print("%s exited" % (e))


c = FSM('a')
c.add_state('a')
c.add_state('b')
c.add_transition('a', 'b', 'ejected', print_enter, print_exit)
print c.pformat()
c.start()
c.process_event("ejected")

print c.current_state
c.process_event("dead")