A state machine in Python

I've written the following code implementing a state machine in Python, and I'd like some feedback on it.

The basic idea is that each state has a set of corresponding actions that trigger a state change. You can register observer functions with states, which are then called whenever the state is entered.

"""
State machine
"""

from typing import Callable, Iterable, Mapping, List, Optional


ObserverFunction = Callable[[str, Optional["State"], "State"], None]


class State:
    """
    Describes a single state
    """
    def __init__(self, name: str, action_map: Mapping[str, str],
                 observers: Optional[Iterable[ObserverFunction]] = None):
        self._name = name
        self._observers: List[ObserverFunction] = list(observers or [])
        self._action_map = action_map


    def __str__(self):
        return self._name


    def register(self, observer: ObserverFunction) -> None:
        """
        Register an observer.
        """
        self._observers.append(observer)


    def enter(self, action: str, prev_state: Optional["State"] = None) -> None:
        """
        Notify all observers that this now is the current state
        """
        for observer in self._observers:
            observer(action, prev_state, self)


    def next_state(self, state_machine: "StateMachine", action: str) -> "State":
        """
        return the new state caused by the action
        """
        if not action in self._action_map:
            raise ValueError(f"State {self} does not support action {action}")
        new_state = state_machine.get_state(self._action_map[action])

        return new_state


    def get_actions(self) -> Iterable[str]:
        """
        return the possible actions on the state
        """
        return self._action_map.keys()


    def is_final(self):
        """
        Return if this state is final, that is, no actions are defined.
        """
        return not self._action_map



class StateMachine:
    """
    Describes a state machine
    """
    def __init__(self, description: Iterable[State], first: str):
        state_names = []
        for state in description:
            name = str(state)
            if name in state_names:
                raise ValueError(f"Duplicate state name: {name}")
            state_names.append(name)

        if not first in state_names:
            raise ValueError(f"No state {first} in {description}")

        self._states = { str(state): state for state in description }
        self._current_state = self._states[first]
        self._current_state.enter("initial state")


    def perform(self, action: str) -> None:
        """
        Perform an action on the current state machine.

        An action potentially changes the state, notifying its observers.
        """
        prev_state = self._current_state
        self._current_state = prev_state.next_state(self, action)
        self._current_state.enter(action, prev_state)


    def get_state(self, name: str) -> State:
        """
        Get state by name
        """
        try:
            return self._states[name]
        except KeyError as error:
            raise ValueError(f"State machine does not contain state {name}") from error


    def get_current_state(self) -> State:
        """
        Get the current state
        """
        return self._current_state


    def has_terminated(self):
        """
        Return True if the state machin has reached a final state
        """
        return self._current_state.is_final()


    def get_all_actions(self):
        """
        Get the set of all actions of all states in the state machine
        """
        all_actions = set()
        for state in self._states.values():
            all_actions.update(state.get_actions())
        return all_actions

The following code demonstrates the intended use:

from statemachine import State, StateMachine


def report(action, old_state, new_state):
    if action == "initial state":
        action = "exist"

    print(f"I was {old_state or 'nothing'}, but you wanted me to {action}, so now I'm {new_state}.") 


def tell_emotion(action, old_state, new_state):
    feelings = {
        'waiting': 'boring',
        'running': 'exhausting',
        'sleeping': 'refreshing' }
    feeling = feelings[str(new_state)]
    print(f"To be {new_state} is {feeling}.")


def final_message(*args):
    print("You monster!")


def main():
    sm = StateMachine([State('waiting',
                             { 'sleep': 'sleeping', 
                               'run': 'running', 
                               'die': 'dead' },
                             [report, tell_emotion]), 
                       State('sleeping',
                             { 'wake up': 'waiting', 
                               'die': 'dead' },
                             [report, tell_emotion]), 
                       State('running',
                             { 'stop': 'waiting', 
                               'die': 'dead' },
                             [report, tell_emotion]), 
                       State('dead', {}, [report, final_message])], 
                      "waiting")
    all_actions = sm.get_all_actions()

    try:
        while not sm.has_terminated(): 
            state = sm.get_current_state()
            print("Available actions:", ", ".join(state.get_actions()))
            try:
                action = input("what should I do? ")
                sm.perform(action)
            except ValueError:
                if action in all_actions:
                    print(f"I can't {action} when I'm {state}!")
                else:
                    print(f"I don't know how to {action}!")
    except EOFError:
        print("Hey, you disappeared!")


if __name__ == '__main__':
    main()

posted over 2 years ago

CC BY-NC-SA 4.0

2y ago by Alexei‭

celtschk‭

771 reputation 8 12 97 2