diff --git a/lib/tfw/yaml_fsm.py b/lib/tfw/yaml_fsm.py
index b4fd45a..9b5f331 100644
--- a/lib/tfw/yaml_fsm.py
+++ b/lib/tfw/yaml_fsm.py
@@ -12,29 +12,29 @@ class YamlFSM(FSMBase):
         self.config = self.parse_config(config_file)
         self.patch_config_callbacks()
         self.setup_states()
-        super(YamlFSM, self).__init__()  # FSMBase.__init__() requires states
+        super().__init__()  # FSMBase.__init__() requires states
         self.setup_transitions()
 
     @staticmethod
     def parse_config(config_file):
         with open(config_file, 'r') as ifile:
-            return yaml.load(ifile)
+            return yaml.safe_load(ifile)
 
     def setup_states(self):
         self.states = [State(**state) for state in self.config['states']]
 
     def setup_transitions(self):
         for transition in self.config['transitions']:
-            self.machine.add_transition(**transition)
+            self.add_transition(**transition)
 
     def patch_config_callbacks(self):
-        topatch = {'on_enter', 'on_exit', 'prepare', 'before', 'after'}
+        topatch = ('on_enter', 'on_exit', 'prepare', 'before', 'after')
 
-        for array in {'states', 'transitions'}:
+        for array in ('states', 'transitions'):
             for json_obj in self.config[array]:
-                for attribute, value in json_obj.items():
-                    if attribute in topatch:
-                        json_obj[attribute] = partial(self.run, value)
+                for key in json_obj:
+                    if key in topatch:
+                        json_obj[key] = partial(self.run, json_obj[key])
 
     @staticmethod
     def run(command, event):