updated error handling for experiment

src/social_norms_trees/ui_wrapper.py

@@ -1,10 +1,10 @@
-import time
 import click
 from datetime import datetime
 import json
 import os
 import uuid
 import py_trees
+import traceback
 
 from social_norms_trees.mutate_tree import move_node, exchange_nodes, remove_node, add_node
 from social_norms_trees.serialize_tree import serialize_tree, deserialize_tree
@@ -92,24 +92,22 @@ def initialize_experiment_record(db, participant_id, origin_tree):
 
 def run_experiment(db, origin_tree, experiment_id, behavior_library):
 
-    # Loop for the actual experiment part, which takes user input to decide which action to take
     print("\nExperiment beginning...\n")
-
-    while(True):
-
-        try:
+
+    try:
+        while True:
             print(py_trees.display.ascii_tree(origin_tree))
             user_choice = click.prompt(
                 "Would you like to perform an action on the behavior tree?",
                 show_choices=True,
                 type=click.Choice(['y', 'n'], case_sensitive=False),
             )
-        
+
             if user_choice == 'y':
                 action = click.prompt(
                     "1. move node\n" +
-                    "2. exchange node\n" + 
-                    "3. remove node\n" + 
+                    "2. exchange node\n" +
+                    "3. remove node\n" +
                     "4. add node\n" +
                     "Please select an action to perform on the behavior tree",
                     type=click.IntRange(min=1, max=4),
@@ -118,33 +116,28 @@ def run_experiment(db, origin_tree, experiment_id, behavior_library):
 
                 if action == 1:
                     origin_tree, action_log = move_node(origin_tree)
-                    db[experiment_id]["action_history"].append(action_log)
                 elif action == 2:
                     origin_tree, action_log = exchange_nodes(origin_tree)
-                    db[experiment_id]["action_history"].append(action_log)
-
                 elif action == 3:
                     origin_tree, action_log = remove_node(origin_tree)
-                    db[experiment_id]["action_history"].append(action_log)
-
                 elif action == 4:
                     origin_tree, action_log = add_node(origin_tree, behavior_library)
-                    db[experiment_id]["action_history"].append(action_log)
 
-                else:
-                    print("Invalid choice, please select a valid number (1, 2, 3, or 4).\n")
+                db[experiment_id]["action_history"].append(action_log)
 
             else:
-                db[experiment_id]["final_behavior_tree"] = serialize_tree(origin_tree)
-                db[experiment_id]["end_date"] = datetime.now().isoformat()
                 break
-        except:
-            print("Error encountered during this action, experiment will now end.")
-            db[experiment_id]["final_behavior_tree"] = serialize_tree(origin_tree)
-            db[experiment_id]["end_date"] = datetime.now().isoformat()
-            break
 
-    return db
+    except Exception:
+        print("\nAn error has occured during the experiment, experiment will now end.")
+        db[experiment_id]["error_log"] = traceback.format_exc()
+
+
+    finally:
+        db[experiment_id]["final_behavior_tree"] = serialize_tree(origin_tree)
+        db[experiment_id]["end_date"] = datetime.now().isoformat()
+        save_db(db, DB_FILE)
+        print("Experiment data has been saved.")
 
 
 
@@ -167,7 +160,7 @@ def main():
 
     participant_id, experiment_id = experiment_setup(db, original_tree)
     db = run_experiment(db, original_tree, experiment_id, behavior_library)
-    save_db(db, DB_FILE)
+    #save_db(db, DB_FILE)
 
     #TODO: define export file, that will be where we export the results to