@@ -234,9 +234,9 @@ def __init__(self, completer, model_completer, docs,
234234 self ._dot_cmd = DotCommandHandler ()
235235 self ._env = os .environ .copy ()
236236 self ._profile = None
237- self .prompt_tokens = u'aws> '
238237 self ._input = input
239238 self ._output = output
239+ self .prompt_tokens = u'aws > '
240240
241241 # These attrs come from the config file.
242242 self .config_obj = None
@@ -272,6 +272,22 @@ def save_config(self):
272272 self .config_section ['theme' ] = self .theme
273273 self .config_obj .write ()
274274
275+ def add_context (self , text ):
276+ self .model_completer .context .append (
277+ text .split ()[0 ].strip ('@' ))
278+ self .model_completer .reset ()
279+ self .prompt_tokens = u'aws ' + ' ' .join (
280+ self .model_completer .context ) + u' > '
281+ self .refresh_cli = True
282+ self .cli .request_redraw ()
283+
284+ def remove_context (self ):
285+ self .model_completer .context .pop ()
286+ self .prompt_tokens = u'aws ' + ' ' .join (
287+ self .model_completer .context ) + u' > '
288+ self .refresh_cli = True
289+ self .cli .request_redraw ()
290+
275291 @property
276292 def cli (self ):
277293 if self ._cli is None or self .refresh_cli :
@@ -283,45 +299,38 @@ def run(self):
283299 while True :
284300 try :
285301 document = self .cli .run (reset_current_buffer = True )
286- if self .model_completer .context and isinstance (self .model_completer .context , list ):
287- text = " " .join (self .model_completer .context ) + " " + document .text
302+ if self .model_completer .context and isinstance (
303+ self .model_completer .context , list ):
304+ text = " " .join (self .model_completer .context ) + \
305+ " " + document .text
306+ original_text = document .text
288307 else :
289308 text = document .text
309+ original_text = text
290310 except InputInterrupt :
291311 pass
292312 except (KeyboardInterrupt , EOFError ):
293313 self .save_config ()
294314 break
295315 else :
296- if text .startswith ('.' ):
316+ if original_text .startswith ('.' ):
297317 # These are special commands (dot commands) that are
298318 # interpreted by the aws-shell directly and typically used
299319 # to modify some type of behavior in the aws-shell.
300320 result = self ._dot_cmd .handle_cmd (text , application = self )
301321 if result is EXIT_REQUESTED :
302322 break
303323 else :
304- if text .startswith ('!' ):
305- # Then run the rest as a normally shell command.
324+ if original_text .startswith ('!' ):
325+ # Then run the rest as a normal shell command.
306326 full_cmd = text [1 :]
307- elif text .startswith ('@' ) and len ( text . split ()) == 1 :
327+ elif original_text .startswith ('@' ):
308328 # Add word as context to completions
309- self .model_completer .context .append (text .split ()[0 ].strip ('@' ))
310- self .model_completer .reset ()
311- self .prompt_tokens = u'aws ' + ' ' .join (self .model_completer .context ) + u' > '
312- self .refresh_cli = True
313- self .cli .request_redraw ()
329+ self .add_context (text )
314330 continue
315- elif 'exit' in text . split () :
331+ elif original_text == 'exit' and self . model_completer . context :
316332 # Remove most recently added context
317- if self .model_completer .context :
318- self .model_completer .context .pop ()
319- if self .model_completer .context :
320- self .prompt_tokens = u'aws ' + ' ' .join (self .model_completer .context ) + u' > '
321- else :
322- self .prompt_tokens = u'aws > '
323- self .refresh_cli = True
324- self .cli .request_redraw ()
333+ self .remove_context ()
325334 continue
326335 else :
327336 full_cmd = 'aws ' + text
0 commit comments