106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
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
|