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