Merge pull request #34 from brown-ccv/fix-add-sequence

fix adding sequence nodes

src/social_norms_trees/mutate_tree.py

@@ -10,7 +10,7 @@
 
 from datetime import datetime
 
-from social_norms_trees.custom_node_library import CustomBehavior
+from social_norms_trees.custom_node_library import CustomBehavior, CustomSequence
 
 
 T = TypeVar("T", bound=py_trees.behaviour.Behaviour)
@@ -153,7 +153,7 @@ def format_parents_with_indices(composite: py_trees.composites.Composite) -> str
     index_strings = []
     i = 0
     for b in iterate_nodes(composite):
-        if b.children:
+        if b.__class__.__name__ == "CustomSequence" or b.__class__.__name__ == "CustomSelector":
             index_strings.append(str(i))
         else:
             index_strings.append("_")
@@ -512,7 +512,7 @@ def exchange_nodes(
 
 def prompt_select_node(behavior_library, text):
 
-    for idx, tree_name in enumerate(behavior_library.behaviors.keys(), 1):
+    for idx, tree_name in enumerate(behavior_library.behavior_from_display_name.keys(), 1):
         print(f"{idx}. {tree_name}")    
 
     choices = [str(i + 1) for i in range(len(behavior_library.behaviors))]
@@ -522,10 +522,9 @@ def prompt_select_node(behavior_library, text):
         show_choices=False
     )
 
-    node_key = list(behavior_library.behaviors.keys())[int(node_index)-1]
-    behavior = behavior_library.behaviors[node_key]
+    node_key = list(behavior_library.behavior_from_display_name.keys())[node_index-1]
 
-    return behavior
+    return behavior_library.behavior_from_display_name[node_key]
 
 
 def add_node(
@@ -542,11 +541,19 @@ def add_node(
 
     behavior = prompt_select_node(behavior_library, f"Which behavior do you want to add?")
 
-    new_node = CustomBehavior(
-                    name=behavior['display_name'],
-                    id_=behavior['id'],
-                    display_name=behavior['display_name']
-                )
+    if behavior['type'] == "Behavior":
+        new_node = CustomBehavior(
+                        name=behavior['display_name'],
+                        id_=behavior['id'],
+                        display_name=behavior['display_name']
+                    )
+
+    elif behavior['type'] == "Sequence":
+        new_node = CustomSequence(
+                        name=behavior['display_name'],
+                        id_=behavior['id'],
+                        display_name=behavior['display_name'],
+                    )
 
     new_parent = prompt_identify_parent_node(
         tree, f"What should its parent be?", display_nodes=True

src/social_norms_trees/serialize_tree.py

@@ -15,12 +15,30 @@ def serialize_node(node):
 
 def deserialize_tree(tree, behavior_library):
     def deserialize_node(node):
+        assert type(node['type'] == str), (
+            f"\nThere was an invalid configuration detected in the inputted behavior tree: "
+            f"Invalid type for node attribute 'type' found for node '{node['name']}'. "
+            f"Please ensure that the 'name' attribute is a string."
+        )
+        assert type(node['name'] == str), (
+            f"\nThere was an invalid configuration detected in the inputted behavior tree: "
+            f"Invalid type for node attribute 'name' found for node '{node['name']}'. "
+            f"Please ensure that the 'name' attribute is a string."
+        )
+
         node_type = node['type']
-        children = [deserialize_node(child) for child in node['children']]
+        assert node_type in ["Sequence", "Selector", "Behavior"], (
+            f"\nThere was an invalid configuration detected in the inputted behavior tree: "
+            f"Invalid node type '{node_type}' found for node '{node['name']}'. "
+            f"Please ensure that all node types are correct and supported."
+        )
 
         behavior = behavior_library.behavior_from_display_name[node['name']]
 
         if node_type == 'Sequence':
+
+            children = [deserialize_node(child) for child in node['children']]
+
             if behavior:
                 return CustomSequence(
                     name=behavior['display_name'],
@@ -30,8 +48,17 @@ def deserialize_node(node):
                 )
             else:
                 raise ValueError(f"Behavior {node['name']} not found in behavior library")
+
+        #TODO: node type Selector 
 
         elif node_type == 'Behavior':
+
+            assert ('children' not in node or len(node['children']) == 0), (
+            f"\nThere was an invalid configuration detected in the inputted behavior tree: "
+            f"Children were detected for Behavior type node '{node['name']}': "
+            f"Behavior nodes should not have any children. Please check the structure of your behavior tree."
+    )
+
             if behavior:
                 return CustomBehavior(
                     name=behavior['display_name'],
@@ -40,5 +67,5 @@ def deserialize_node(node):
                 )
             else:
                 raise ValueError(f"Behavior {node['name']} not found in behavior library")
-            
+
     return deserialize_node(tree)

src/social_norms_trees/ui_wrapper.py

@@ -105,7 +105,7 @@ def run_experiment(db, origin_tree, experiment_id, behavior_library):
                 show_choices=True,
                 type=click.Choice(['y', 'n'], case_sensitive=False),
             )
-
+            
             if user_choice == 'y':
                 action = click.prompt(
                     "1. move node\n" +
@@ -138,7 +138,7 @@ def run_experiment(db, origin_tree, experiment_id, behavior_library):
             #user_choice == "n", end simulation run
             else:
                 break
-        
+
     except Exception:
         print("\nAn error has occured during the experiment, the experiment will now end.")
         db[experiment_id]["error_log"] = traceback.format_exc()
@@ -164,9 +164,9 @@ def main():
 
     RESOURCES_FILE = "resources.json"
     original_tree, behavior_library, context_paragraph = load_resources(RESOURCES_FILE)
-
     print(f"\nContext of this experiment: {context_paragraph}")
 
+
     participant_id, experiment_id = experiment_setup(db, original_tree)
     db = run_experiment(db, original_tree, experiment_id, behavior_library)
     save_db(db, DB_FILE)