baseimage-tutorial-framework/lib/tfw/fsm/yaml_fsm.py

from subprocess import Popen, run
from functools import partial, singledispatch
from contextlib import suppress

import yaml
import jinja2
from transitions import State

from .fsm_base import FSMBase


class YamlFSM(FSMBase):
    """
    This is a state machine capable of building itself from a YAML config file.
    """
    def __init__(self, config_file, jinja2_variables=None):
        """
        :param config_file: path of the YAML file
        :param jinja2_variables: dict containing jinja2 variables
                                 or str with filename of YAML file to
                                 parse and use as dict.
                                 jinja2 support is disabled if this is None
        """
        self.config = ConfigParser(config_file, jinja2_variables).config
        self.setup_states()
        super().__init__()  # FSMBase.__init__() requires states
        self.setup_transitions()

    def setup_states(self):
        self.for_config_states_and_transitions_do(self.wrap_callbacks_with_subprocess_call)
        self.states = [State(**state) for state in self.config['states']]

    def setup_transitions(self):
        self.for_config_states_and_transitions_do(self.subscribe_and_remove_predicates)
        for transition in self.config['transitions']:
            self.add_transition(**transition)

    def for_config_states_and_transitions_do(self, what):
        for array in ('states', 'transitions'):
            for json_obj in self.config[array]:
                what(json_obj)

    @staticmethod
    def wrap_callbacks_with_subprocess_call(json_obj):
        topatch = ('on_enter', 'on_exit', 'prepare', 'before', 'after')
        for key in json_obj:
            if key in topatch:
                json_obj[key] = partial(run_command_async, json_obj[key])

    def subscribe_and_remove_predicates(self, json_obj):
        if 'predicates' in json_obj:
            for predicate in json_obj['predicates']:
                self.subscribe_predicate(
                    json_obj['trigger'],
                    partial(
                        command_statuscode_is_zero,
                        predicate
                    )
                )

        with suppress(KeyError):
            json_obj.pop('predicates')


def run_command_async(command, _):
    Popen(command, shell=True)


def command_statuscode_is_zero(command):
    return run(command, shell=True).returncode == 0


class ConfigParser:
    def __init__(self, config_file, jinja2_variables):
        self.read_variables = singledispatch(self._read_variables)
        self.read_variables.register(dict, self._read_variables_dict)
        self.read_variables.register(str, self._read_variables_str)

        self.config = self.parse_config(config_file, jinja2_variables)

    def parse_config(self, config_file, jinja2_variables):
        config_string = self.read_file(config_file)
        if jinja2_variables is not None:
            variables = self.read_variables(jinja2_variables)
            template = jinja2.Environment(loader=jinja2.BaseLoader).from_string(config_string)
            config_string = template.render(**variables)
        return yaml.safe_load(config_string)

    @staticmethod
    def read_file(filename):
        with open(filename, 'r') as ifile:
            return ifile.read()

    @staticmethod
    def _read_variables(variables):
        raise TypeError(f'Invalid variables type {type(variables)}')

    @staticmethod
    def _read_variables_str(variables):
        with open(variables, 'r') as ifile:
            return yaml.safe_load(ifile)

    @staticmethod
    def _read_variables_dict(variables):
        return variables