Implemented stage history

Author: joguSD

awsshell/wizard.py

@@ -1,11 +1,15 @@
+import sys
+import copy
 import json
 import jmespath
 
 import botocore.session
 from botocore import xform_name
 
 from awsshell.resource import index
-from awsshell.interaction import InteractionLoader
+from awsshell.interaction import InteractionLoader, SimpleSelect
+
+from prompt_toolkit.shortcuts import confirm
 
 
 class WizardException(Exception):
@@ -118,6 +122,13 @@ def _load_stages(self, stages):
         for stage in stages:
             self.stages[stage.name] = stage
 
+    def _select_stage_prompt(self):
+        stages = [s.name for (s, _) in self.env._stage_history]
+        selector = SimpleSelect({}, 'Select a Stage to return to: ')
+        current_stage = selector.execute(stages)
+        self.env.pop_stages(stages.index(current_stage))
+        return current_stage
+
     def execute(self):
         """Run the wizard. Execute Stages until a final stage is reached.
 
@@ -128,8 +139,18 @@ def execute(self):
             stage = self.stages.get(current_stage)
             if not stage:
                 raise WizardException('Stage not found: %s' % current_stage)
-            stage.execute()
-            current_stage = stage.get_next_stage()
+            try:
+                stage.execute()
+                current_stage = stage.get_next_stage()
+            except EOFError:
+                current_stage = self._select_stage_prompt()
+            except Exception as e:
+                sys.stdout.write('{0}\n'.format(e))
+                sys.stdout.flush()
+                if confirm('Select a previous stage? (y/n) '):
+                    current_stage = self._select_stage_prompt()
+                else:
+                    raise
 
 
 class Stage(object):
@@ -245,6 +266,7 @@ def execute(self):
         2) Perform Interaction on retrieved data.
         3) Perform Resolution to store data in the environment.
         """
+        self._env.push_stage(self)
         retrieved_options = self._handle_retrieval()
         selected_data = self._handle_interaction(retrieved_options)
         self._handle_resolution(selected_data)
@@ -253,8 +275,9 @@ def execute(self):
 class Environment(object):
     """Store vars into a dict and retrives them with JMESPath queries."""
 
-    def __init__(self):
+    def __init__(self, stage_history=[]):
         self._variables = {}
+        self._stage_history = stage_history
 
     def __str__(self):
         return json.dumps(self._variables, indent=4, sort_keys=True)
@@ -291,6 +314,14 @@ def resolve_parameters(self, keys):
         :rtype: dict
         :return: The dict of with all of the paths resolved to their values.
         """
+        resolved_dict = {}
         for key in keys:
-            keys[key] = self.retrieve(keys[key])
-        return keys
+            resolved_dict[key] = self.retrieve(keys[key])
+        return resolved_dict
+
+    def push_stage(self, stage):
+        self._stage_history.append((stage, copy.deepcopy(self._variables)))
+
+    def pop_stages(self, stage_index):
+        self._variables = self._stage_history[stage_index][1]
+        self._stage_history = self._stage_history[:stage_index]