Merge pull request awslabs#81 from awslabs/lexer

Implements awslabs#27: Add lexer/syntax highlighting.

Author: donnemartin

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)
 

awsshell/index/completion.py

@@ -9,16 +9,31 @@
 
 """
 import os
+import json
 
 from awsshell.utils import FSLayer, FileReadError, build_config_file_path
+from awsshell import utils
 
 
 class IndexLoadError(Exception):
     """Raised when an index could not be loaded."""
 
 
 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,13 +45,18 @@ 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.
 
         :type version_string: str
         :param version_string: The AWS CLI version, e.g "1.9.2".
 
+        :raises: :class:`IndexLoadError <exceptions.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)

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),
+        ]
+    }

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 = [

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'])