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

# Copyright (C) 2018 Avatao.com Innovative Learning Kft.
# All Rights Reserved. See LICENSE file for details.

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

import yaml
import jinja2
from transitions import State

from tfw.fsm.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