Stateful commands + small cache and performance changes

awsshell/app.py

@@ -236,6 +236,7 @@ def __init__(self, completer, model_completer, docs,
         self._profile = None
         self._input = input
         self._output = output
+        self.prompt_tokens = u'aws > '
 
         # These attrs come from the config file.
         self.config_obj = None
@@ -271,6 +272,22 @@ def save_config(self):
         self.config_section['theme'] = self.theme
         self.config_obj.write()
 
+    def add_context(self, text):
+        self.model_completer.context.append(
+            text.split()[0].strip('@'))
+        self.model_completer.reset()
+        self.prompt_tokens = u'aws ' + ' '.join(
+            self.model_completer.context) + u' > '
+        self.refresh_cli = True
+        self.cli.request_redraw()
+
+    def remove_context(self):
+        self.model_completer.context.pop()
+        self.prompt_tokens = u'aws ' + ' '.join(
+            self.model_completer.context) + u' > '
+        self.refresh_cli = True
+        self.cli.request_redraw()
+
     @property
     def cli(self):
         if self._cli is None or self.refresh_cli:
@@ -282,24 +299,40 @@ def run(self):
         while True:
             try:
                 document = self.cli.run(reset_current_buffer=True)
-                text = document.text
+                if self.model_completer.context and isinstance(
+                        self.model_completer.context, list):
+                    text = " ".join(self.model_completer.context) + \
+                        " " + document.text
+                    original_text = document.text
+                else:
+                    text = document.text
+                    original_text = text
             except InputInterrupt:
                 pass
             except (KeyboardInterrupt, EOFError):
                 self.save_config()
                 break
             else:
-                if text.startswith('.'):
+                if original_text.startswith('.'):
                     # These are special commands (dot commands) that are
                     # interpreted by the aws-shell directly and typically used
                     # to modify some type of behavior in the aws-shell.
                     result = self._dot_cmd.handle_cmd(text, application=self)
                     if result is EXIT_REQUESTED:
                         break
                 else:
-                    if text.startswith('!'):
-                        # Then run the rest as a normally shell command.
+                    if original_text.startswith('!'):
+                        # Then run the rest as a normal shell command.
                         full_cmd = text[1:]
+                    elif original_text.startswith('@'):
+                        # Add word as context to completions
+                        self.add_context(text)
+                        continue
+                    elif original_text == 'exit' and\
+                            self.model_completer.context:
+                        # Remove most recently added context
+                        self.remove_context()
+                        continue
                     else:
                         full_cmd = 'aws ' + text
                         self.history.append(full_cmd)
@@ -331,7 +364,7 @@ def create_layout(self, display_completions_in_columns, toolbar):
         if self.config_section['theme'] == 'none':
             lexer = None
         return create_default_layout(
-            self, u'aws> ', lexer=lexer, reserve_space_for_menu=True,
+            self, self.prompt_tokens, lexer=lexer, reserve_space_for_menu=True,
             display_completions_in_columns=display_completions_in_columns,
             get_bottom_toolbar_tokens=toolbar.handler)
 

awsshell/autocomplete.py

@@ -24,6 +24,8 @@ def __init__(self, index_data, match_fuzzy=True):
         # This will get populated as a command is completed.
         self.cmd_path = [self._current_name]
         self.match_fuzzy = match_fuzzy
+        self.context = []
+        self._cache_all_args = []
 
     @property
     def global_arg_metadata(self):
@@ -42,6 +44,16 @@ def reset(self):
         self._last_position = 0
         self.last_option = ''
         self.cmd_path = [self._current_name]
+        for context in self.context:
+            next_command = self._current['children'].get(context)
+            if not next_command:
+                self.context.remove(context)
+                self.reset()
+                return
+            self._current = next_command
+            self._current_name = context
+            self.cmd_path.append(self._current_name)
+        self._cache_all_args = []
 
     def autocomplete(self, line):
         """Given a line, return a list of suggestions."""
@@ -57,7 +69,10 @@ def autocomplete(self, line):
             return self._handle_backspace()
         elif not line:
             return []
-        elif current_length != self._last_position + 1:
+        elif (self._last_position == 0 and
+                current_length > 2) or (
+                current_length != self._last_position + 1
+                and '--' in line):
             return self._complete_from_full_parse()
 
         # This position is important.  We only update the _last_position
@@ -75,6 +90,7 @@ def autocomplete(self, line):
             # this as self.last_arg
             self.last_option = last_word
         if line[-1] == ' ':
+            self._cache_all_args = []
             # At this point the user has autocompleted a command
             # or an argument and has hit space.  If they've
             # just completed a command, we need to change the
@@ -101,14 +117,12 @@ def autocomplete(self, line):
             # in either of the above two cases.
             return self._current['commands'][:]
         elif last_word.startswith('-'):
-            # TODO: cache this for the duration of the current context.
-            # We don't need to recompute this until the args are
-            # different.
-            all_args = self._get_all_args()
+            if not self._cache_all_args:
+                self._cache_all_args = self._get_all_args()
             if self.match_fuzzy:
-                return fuzzy_search(last_word, all_args)
+                return fuzzy_search(last_word, self._cache_all_args)
             else:
-                return substring_search(last_word, all_args)
+                return substring_search(last_word, self._cache_all_args)
         if self.match_fuzzy:
             return fuzzy_search(last_word, self._current['commands'])
         else:

tests/unit/test_autocomplete.py

@@ -380,3 +380,31 @@ def test_global_arg_metadata_property(index_data):
     }
     completer = AWSCLIModelCompleter(index_data)
     assert '--global1' in completer.global_arg_metadata
+
+def test_add_context_changes_context(index_data):
+    index_data['aws']['commands'] = ['ec2']
+    index_data['aws']['children'] = {
+        'ec2': {
+            'commands': ['create-tags'],
+            'argument_metadata': {},
+            'arguments': [],
+            'children': {
+                'create-tags': {
+                    'commands': [],
+                    'argument_metadata': {
+                        '--resources': {'example': '', 'minidoc': 'foo'},
+                    },
+                    'arguments': ['--resources'],
+                    'children': {},
+                }
+            }
+        }
+    }
+    completer = AWSCLIModelCompleter(index_data)
+    completer.reset()
+    completer.autocomplete('c')
+    assert completer.autocomplete('c') == ['ec2']
+    completer.context = ['ec2']
+    completer.reset()
+    completer.autocomplete('c')
+    assert completer.autocomplete('c') == ['create-tags']