baseimage-tutorial-framework/tfw/fsm/fsm_base.py

import logging
from collections import defaultdict
from datetime import datetime
from contextlib import suppress

from transitions import Machine, MachineError

from tfw.internals.callback_mixin import CallbackMixin

LOG = logging.getLogger(__name__)


class FSMBase(Machine, CallbackMixin):
    """
    A general FSM base class you can inherit from to track user progress.
    See linear_fsm.py for an example use-case.
    TFW uses the transitions library for state machines, please refer to their
    documentation for more information on creating your own machines:
    https://github.com/pytransitions/transitions
    """
    states, transitions = [], []

    def __init__(self, initial=None, accepted_states=None):
        """
        :param initial: which state to begin with, defaults to the last one
        :param accepted_states: list of states in which the challenge should be
                                considered successfully completed
        """
        self.accepted_states = accepted_states or [self.states[-1].name]
        self.trigger_predicates = defaultdict(list)
        self.event_log = []

        Machine.__init__(
            self,
            states=self.states,
            transitions=self.transitions,
            initial=initial or self.states[0],
            send_event=True,
            ignore_invalid_triggers=True,
            after_state_change='execute_callbacks'
        )

    def execute_callbacks(self, event_data):
        self._execute_callbacks(event_data.kwargs)

    def is_solved(self):
        return self.state in self.accepted_states  # pylint: disable=no-member

    def subscribe_predicate(self, trigger, *predicates):
        self.trigger_predicates[trigger].extend(predicates)

    def unsubscribe_predicate(self, trigger, *predicates):
        self.trigger_predicates[trigger] = [
            predicate
            for predicate in self.trigger_predicates[trigger]
            not in predicates
        ]

    def step(self, trigger):
        predicate_results = (
            predicate()
            for predicate in self.trigger_predicates[trigger]
        )

        if all(predicate_results):
            with suppress(AttributeError, MachineError):
                from_state = self.state
                if self.trigger(trigger):
                    self.update_event_log(from_state, trigger)
                    return True
            LOG.debug('FSM failed to execute nonexistent trigger: "%s"', trigger)

    def update_event_log(self, from_state, trigger):
        self.event_log.append({
            'from_state': from_state,
            'to_state': self.state,
            'trigger': trigger,
            'timestamp': datetime.utcnow().isoformat()
        })

    @property
    def in_accepted_state(self):
        return self.state in self.accepted_states