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

96 lines
3.2 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 tfw import FSMBase
class YamlFSM(FSMBase):
def __init__(self, config_file, jinja2_variables=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