Python – State Machine Code

'''
    COMP-E, reposted as __nero on activestate
    Monday, May 08, 2011
'''

import inspect
import sys
import re

class StateMachineFoo(object):
    '''
        This is a template for a finite state machine that is dynamically built based
        on the function names and docstring for the given functions
    '''

    def __init__(self):
        # find all the states based on the regex
        m = re.compile(r'_state.*',re.I)

        # return a list of matches
        state_names = filter(m.search,dir(self))

        # get the function pointers to all of our states
        state_ptrs = [getattr(self,state) for state in state_names]

        # from the function pointers, suck up our dockstring and get our transitions
        dictm = re.compile(r"{.*}",re.S)
        state_trans = [eval(dictm.search(getattr(state, 'func_doc', '')).group(0)) for state in state_ptrs]

        # use zip to create a tuple containing our state_name/state_transition pairings, and add
        # them to the dictionary as key/value pairs
        self.state_machine = dict(zip(state_names,state_trans))

    def start(self, start_state):
        '''
            you should specify your beginning state
        '''
        if type('') == type(start_state):
            getattr(self,start_state)()
        else:
            self._stateA()

    def stop(self):
        '''
            A state to exit nicely from
        '''
        return

    def _stateA(self):
        '''
            {'SUCCESS': self._stateB,'FAIL':self.error}
        '''
        error = False
        fname = inspect.stack()[0][3]

        print "I'm in '%s'" % fname

        if error:
            self.state_machine[fname]['FAIL']()
        else:
            self.state_machine[fname]['SUCCESS']()

    def _stateB(self):
        '''
            {'SUCCESS': self._stateD,'FAIL':self.error}
        '''
        error = False
        fname = inspect.stack()[0][3]

        print "I'm in '%s'" % fname

        if error:
            self.state_machine[fname]['FAIL']()
        else:
            self.state_machine[fname]['SUCCESS']()

    def _stateC(self):
        '''
            {'SUCCESS': self._stateA,'FAIL':self.error}
        '''
        error = True
        fname = inspect.stack()[0][3]

        print "I'm in '%s'" % fname

        if error:
            self.state_machine[fname]['FAIL']()
        else:
            self.state_machine[fname]['SUCCESS']()

    def _stateD(self):
        '''
            {'SUCCESS': self._stateE,'FAIL':self.error}
        '''
        error = False
        fname = inspect.stack()[0][3]

        print "I'm in '%s'" % fname

        if error:
            self.state_machine[fname]['FAIL']()
        else:
            self.state_machine[fname]['SUCCESS']()

    def _stateE(self):
        '''
            {'SUCCESS': self._stateC,'FAIL':self.error}
        '''
        error = False
        fname = inspect.stack()[0][3]

        print "I'm in '%s'" % fname

        if error:
            self.state_machine[fname]['FAIL']()
        else:
            self.state_machine[fname]['SUCCESS']()

    def error(self):
        print "OH NO, AN ERROR!!! "
        sys.exit(-1)

if __name__ == '__main__':

    my_self     = StateMachineFoo()
    start_state = None

    # this is a really bad way to do this, probably want python argparse module, but for
    # demo purposes, this should allow you to change your start state from the command line
    # the simple/dirty way

    if len(sys.argv) > 1:
        start_state = sys.argv[1]

    my_self.start(start_state)