baseimage-tutorial-framework/lib/tfw/fsm/yaml_fsm.py
2018-07-31 13:34:00 +02:00

109 lines
3.7 KiB
Python

# 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