PTA: Load state power functions from YAML if available

1 files changed, 5 insertions, 1 deletions

diff --git a/lib/automata.py b/lib/automata.py
index 552b482..d977164 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -102,7 +102,7 @@ class PTAAttribute:
 
     @classmethod
     def from_json_maybe(cls, json_wrapped: dict, attribute: str, parameters: dict):
-        if attribute in json_wrapped:
+        if type(json_wrapped) is dict and attribute in json_wrapped:
             return cls.from_json(json_wrapped[attribute], parameters)
         return cls()
 
@@ -653,6 +653,10 @@ class PTA:
 
         pta = cls(**kwargs)
 
+        if 'state' in yaml_input:
+            for state_name, state in yaml_input['state'].items():
+                pta.add_state(state_name, power=PTAAttribute.from_json_maybe(state, 'power', pta.parameters))
+
         for trans_name in sorted(yaml_input['transition'].keys()):
             kwargs = dict()
             transition = yaml_input['transition'][trans_name]