diff --git a/awsshell/app.py b/awsshell/app.py index 8340dfe..218e2b0 100644 --- a/awsshell/app.py +++ b/awsshell/app.py @@ -239,8 +239,12 @@ def stop_input_and_refresh_cli(self): raise InputInterrupt def create_layout(self, display_completions_in_columns, toolbar): + from awsshell.lexer import ShellLexer + lexer = ShellLexer + if self.config_section['theme'] == 'none': + lexer = None return create_default_layout( - self, u'aws> ', reserve_space_for_menu=True, + self, u'aws> ', lexer=lexer, reserve_space_for_menu=True, display_completions_in_columns=display_completions_in_columns, get_bottom_toolbar_tokens=toolbar.handler) diff --git a/awsshell/index/completion.py b/awsshell/index/completion.py index eaa7374..5737a09 100644 --- a/awsshell/index/completion.py +++ b/awsshell/index/completion.py @@ -9,8 +9,10 @@ """ import os +import json from awsshell.utils import FSLayer, FileReadError, build_config_file_path +from awsshell import utils class IndexLoadError(Exception): @@ -18,7 +20,20 @@ class IndexLoadError(Exception): class CompletionIndex(object): - """Handles working with the local commmand completion index.""" + """Handles working with the local commmand completion index. + + :type commands: list + :param commands: ec2, s3, elb... + + :type subcommands: list + :param subcommands: start-instances, stop-instances, terminate-instances... + + :type global_opts: list + :param global_opts: --profile, --region, --output... + + :type args_opts: set, to filter out duplicates + :param args_opts: ec2 start-instances: --instance-ids, --dry-run... + """ # The completion index can read/write to a cache dir # so that it doesn't have to recompute the completion cache @@ -30,6 +45,10 @@ def __init__(self, cache_dir=DEFAULT_CACHE_DIR, fslayer=None): if fslayer is None: fslayer = FSLayer() self._fslayer = fslayer + self.commands = [] + self.subcommands = [] + self.global_opts = [] + self.args_opts = set() def load_index(self, version_string): """Load the completion index for a given CLI version. @@ -37,6 +56,7 @@ def load_index(self, version_string): :type version_string: str :param version_string: The AWS CLI version, e.g "1.9.2". + :raises: :class:`IndexLoadError ` """ filename = self._filename_for_version(version_string) try: @@ -48,3 +68,35 @@ def load_index(self, version_string): def _filename_for_version(self, version_string): return os.path.join( self._cache_dir, 'completions-%s.json' % version_string) + + def load_completions(self): + """Loads completions from the completion index. + + Updates the following attributes: + * commands + * subcommands + * global_opts + * args_opts + """ + try: + index_str = self.load_index(utils.AWSCLI_VERSION) + except IndexLoadError: + return + index_str = self.load_index(utils.AWSCLI_VERSION) + index_data = json.loads(index_str) + index_root = index_data['aws'] + # ec2, s3, elb... + self.commands = index_root['commands'] + # --profile, --region, --output... + self.global_opts = index_root['arguments'] + for command in self.commands: + # ec2: start-instances, stop-instances, terminate-instances... + subcommands_current = index_root['children'] \ + .get(command)['commands'] + self.subcommands.extend(subcommands_current) + for subcommand_current in subcommands_current: + # start-instances: --instance-ids, --dry-run... + args_opts_current = index_root['children'] \ + .get(command)['children'] \ + .get(subcommand_current)['arguments'] + self.args_opts.update(args_opts_current) diff --git a/awsshell/lexer.py b/awsshell/lexer.py new file mode 100644 index 0000000..287092a --- /dev/null +++ b/awsshell/lexer.py @@ -0,0 +1,65 @@ +# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +import ast +import os + +from pygments.lexer import RegexLexer +from pygments.lexer import words +from pygments.token import Keyword, Literal, Name, Operator, Text + +from awsshell.index.completion import CompletionIndex + + +class ShellLexer(RegexLexer): + """Provides highlighting for commands, subcommands, arguments, and options. + + :type completion_index: :class:`CompletionIndex` + :param completion_index: Completion index used to determine commands, + subcommands, arguments, and options for highlighting. + + :type tokens: dict + :param tokens: A dict of (`pygments.lexer`, `pygments.token`) used for + pygments highlighting. + """ + completion_index = CompletionIndex() + completion_index.load_completions() + tokens = { + 'root': [ + # ec2, s3, elb... + (words( + tuple(completion_index.commands), + prefix=r'\b', + suffix=r'\b'), + Literal.String), + # describe-instances + (words( + tuple(completion_index.subcommands), + prefix=r'\b', + suffix=r'\b'), + Name.Class), + # --instance-ids + (words( + tuple(list(completion_index.args_opts)), + prefix=r'', + suffix=r'\b'), + Keyword.Declaration), + # --profile + (words( + tuple(completion_index.global_opts), + prefix=r'', + suffix=r'\b'), + Operator.Word), + # Everything else + (r'.*\n', Text), + ] + } diff --git a/awsshell/ui.py b/awsshell/ui.py index c491ed9..657c9d3 100644 --- a/awsshell/ui.py +++ b/awsshell/ui.py @@ -16,7 +16,9 @@ from prompt_toolkit.layout.toolbars import ValidationToolbar, \ SystemToolbar, ArgToolbar, SearchToolbar from prompt_toolkit.layout.utils import explode_tokens +from prompt_toolkit.layout.lexers import PygmentsLexer from pygments.token import Token +from pygments.lexer import Lexer from awsshell.compat import text_type @@ -65,6 +67,16 @@ def create_default_layout(app, message='', get_prompt_tokens_1, get_prompt_tokens_2 = _split_multiline_prompt( get_prompt_tokens) + + # `lexer` is supposed to be a `Lexer` instance. But if a Pygments lexer + # class is given, turn it into a PygmentsLexer. (Important for + # backwards-compatibility.) + try: + if issubclass(lexer, Lexer): + lexer = PygmentsLexer(lexer) + except TypeError: # Happens when lexer is `None` or an instance of something else. + pass + # Create processors list. # (DefaultPrompt should always be at the end.) input_processors = [ diff --git a/tests/unit/test_load_completions.py b/tests/unit/test_load_completions.py new file mode 100644 index 0000000..ae5086b --- /dev/null +++ b/tests/unit/test_load_completions.py @@ -0,0 +1,49 @@ +# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +import unittest + +from awsshell.index.completion import CompletionIndex + + +class LoadCompletionsTest(unittest.TestCase): + + def setUp(self): + self.completion_index = CompletionIndex() + # This would probably be cleaner with a pytest.fixture like + # test_completions.index_data + DATA = ( + '{"aws": ' + '{"commands": ["devicefarm", "foo"], ' + '"arguments": ["--debug", "--endpoint-url"], ' + '"children": {"devicefarm": ' + '{"commands": ["create-device-pool"], ' + '"children": {"create-device-pool": ' + '{"commands": [], ' + '"arguments": ["--project-arn", "--name"]}}}, ' + '"foo": ' + '{"commands": ["bar"], ' + '"children": {"bar": ' + '{"commands": [], "arguments": ["--baz"]}}}}}}' + ) + self.completion_index.load_index = lambda x: DATA + self.completion_index.load_completions() + + def test_load_completions(self): + assert self.completion_index.commands == [ + 'devicefarm', 'foo'] + assert self.completion_index.subcommands == [ + 'create-device-pool', 'bar'] + assert self.completion_index.global_opts == [ + '--debug', '--endpoint-url'] + assert self.completion_index.args_opts == set([ + '--project-arn', '--name', '--baz'])