Tab completion for Alias (#127)

* Tab completion for Alias

* Use tab completion table approach

* Address PR comments

* Change 2.0.31.dev0

* Fix CI error

src/alias/azext_alias/__init__.py

@@ -3,27 +3,28 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import timeit
-
-from knack.log import get_logger
-
 from azure.cli.core import AzCommandsLoader
 from azure.cli.core.decorators import Completer
-from azure.cli.core.commands.events import EVENT_INVOKER_PRE_CMD_TBL_TRUNCATE
-from azext_alias.alias import GLOBAL_ALIAS_PATH, AliasManager
-from azext_alias.util import get_config_parser, is_alias_create_command, cache_reserved_commands
-from azext_alias._const import DEBUG_MSG_WITH_TIMING
-from azext_alias._validators import process_alias_create_namespace
-from azext_alias import telemetry
+from azure.cli.core.commands.events import EVENT_INVOKER_PRE_CMD_TBL_TRUNCATE, EVENT_INVOKER_ON_TAB_COMPLETION
+from azure.cli.command_modules.interactive.events import (
+    EVENT_INTERACTIVE_PRE_COMPLETER_TEXT_PARSING,
+    EVENT_INTERACTIVE_POST_SUB_TREE_CREATE
+)
 from azext_alias import _help  # pylint: disable=unused-import
+from azext_alias.hooks import (
+    alias_event_handler,
+    enable_aliases_autocomplete,
+    transform_cur_commands_interactive,
+    enable_aliases_autocomplete_interactive
+)
+from azext_alias.util import get_alias_table
+from azext_alias._validators import process_alias_create_namespace
 
-logger = get_logger(__name__)
-"""
-We don't have access to load_cmd_tbl_func in custom.py (need the entire command table
-for alias and command validation when the user invokes alias create).
-This cache saves the entire command table globally so custom.py can have access to it.
-Alter this cache through cache_reserved_commands(load_cmd_tbl_func) in util.py
-"""
+
+# We don't have access to load_cmd_tbl_func in custom.py (need the entire command table
+# for alias and command validation when the user invokes alias create).
+# This cache saves the entire command table globally so custom.py can have access to it.
+# Alter this cache through cache_reserved_commands(load_cmd_tbl_func) in util.py
 cached_reserved_commands = []
 
 
@@ -35,6 +36,9 @@ def __init__(self, cli_ctx=None):
         super(AliasExtCommandLoader, self).__init__(cli_ctx=cli_ctx,
                                                     custom_command_type=custom_command_type)
         self.cli_ctx.register_event(EVENT_INVOKER_PRE_CMD_TBL_TRUNCATE, alias_event_handler)
+        self.cli_ctx.register_event(EVENT_INVOKER_ON_TAB_COMPLETION, enable_aliases_autocomplete)
+        self.cli_ctx.register_event(EVENT_INTERACTIVE_PRE_COMPLETER_TEXT_PARSING, transform_cur_commands_interactive)
+        self.cli_ctx.register_event(EVENT_INTERACTIVE_POST_SUB_TREE_CREATE, enable_aliases_autocomplete_interactive)
 
     def load_command_table(self, _):
         with self.command_group('alias') as g:
@@ -59,41 +63,7 @@ def get_alias_completer(cmd, prefix, namespace, **kwargs):  # pylint: disable=un
     """
     An argument completer for alias name.
     """
-    try:
-        alias_table = get_config_parser()
-        alias_table.read(GLOBAL_ALIAS_PATH)
-        return alias_table.sections()
-    except Exception:  # pylint: disable=broad-except
-        return []
-
-
-def alias_event_handler(_, **kwargs):
-    """
-    An event handler for alias transformation when EVENT_INVOKER_PRE_TRUNCATE_CMD_TBL event is invoked
-    """
-    try:
-        telemetry.start()
-
-        start_time = timeit.default_timer()
-        args = kwargs.get('args')
-        alias_manager = AliasManager(**kwargs)
-
-        # [:] will keep the reference of the original args
-        args[:] = alias_manager.transform(args)
-
-        if is_alias_create_command(args):
-            load_cmd_tbl_func = kwargs.get('load_cmd_tbl_func', lambda _: {})
-            cache_reserved_commands(load_cmd_tbl_func)
-
-        elapsed_time = (timeit.default_timer() - start_time) * 1000
-        logger.debug(DEBUG_MSG_WITH_TIMING, args, elapsed_time)
-
-        telemetry.set_execution_time(round(elapsed_time, 2))
-    except Exception as client_exception:  # pylint: disable=broad-except
-        telemetry.set_exception(client_exception)
-        raise
-    finally:
-        telemetry.conclude()
+    return get_alias_table().sections()
 
 
 COMMAND_LOADER_CLS = AliasExtCommandLoader

src/alias/azext_alias/_const.py

@@ -3,12 +3,16 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+import os
+
 from azure.cli.core._environment import get_config_dir
 
 GLOBAL_CONFIG_DIR = get_config_dir()
 ALIAS_FILE_NAME = 'alias'
 ALIAS_HASH_FILE_NAME = 'alias.sha1'
 COLLIDED_ALIAS_FILE_NAME = 'collided_alias'
+ALIAS_TAB_COMP_TABLE_FILE_NAME = 'alias_tab_completion'
+GLOBAL_ALIAS_TAB_COMP_TABLE_PATH = os.path.join(GLOBAL_CONFIG_DIR, ALIAS_TAB_COMP_TABLE_FILE_NAME)
 COLLISION_CHECK_LEVEL_DEPTH = 5
 
 INSUFFICIENT_POS_ARG_ERROR = 'alias: "{}" takes exactly {} positional argument{} ({} given)'

src/alias/azext_alias/_help.py

@@ -16,23 +16,26 @@
     type: command
     short-summary: Create an alias.
     examples:
-        - name: Create a simple alias.
-          text: >
-            az alias create --name rg --command group\n
+        - name: Create simple alias commands.
+          text: |
+            az alias create --name rg --command group
+
             az alias create --name ls --command list
         - name: Create a complex alias.
-          text: >
+          text: |
             az alias create --name list-vm --command 'vm list --resource-group myResourceGroup'
 
-        - name: Create an alias with positional arguments.
-          text: >
-            az alias create --name 'list-vm {{ resource_group }}' --command 'vm list --resource-group {{ resource_group }}'
-
-        - name: Create an alias with positional arguments and additional string processing.
-          text: >
-            az alias create --name 'storage-ls {{ url }}' --command 'storage blob list \n
-            --account-name {{ url.replace("https://", "").split(".")[0] }}\n
-            --container-name {{ url.replace("https://", "").split("/")[1] }}'
+        - name: Create an alias command with arguments.
+          text: |
+            az alias create --name 'list-vm {{ resource_group }}' \\
+              --command 'vm list --resource-group {{ resource_group }}'
+
+        - name: Process arguments using Jinja2 templates.
+          text: |
+            az alias create --name 'storage-ls {{ url }}' \\
+              --command 'storage blob list
+                --account-name {{ url.replace("https://", "").split(".")[0] }}
+                --container-name {{ url.replace("https://", "").split("/")[1] }}'
 """
 
 

src/alias/azext_alias/_validators.py

@@ -107,13 +107,13 @@ def _validate_alias_command_level(alias, command):
         alias: The name of the alias.
         command: The command that the alias points to.
     """
-    alias_collision_table = AliasManager.build_collision_table([alias], azext_alias.cached_reserved_commands)
+    alias_collision_table = AliasManager.build_collision_table([alias])
 
     # Alias is not a reserved command, so it can point to any command
     if not alias_collision_table:
         return
 
-    command_collision_table = AliasManager.build_collision_table([command], azext_alias.cached_reserved_commands)
+    command_collision_table = AliasManager.build_collision_table([command])
     alias_collision_levels = alias_collision_table.get(alias.split()[0], [])
     command_collision_levels = command_collision_table.get(command.split()[0], [])
 

src/alias/azext_alias/alias.py

@@ -25,7 +25,12 @@
     POS_ARG_DEBUG_MSG
 )
 from azext_alias.argument import build_pos_args_table, render_template
-from azext_alias.util import is_alias_create_command, cache_reserved_commands, get_config_parser
+from azext_alias.util import (
+    is_alias_create_command,
+    cache_reserved_commands,
+    get_config_parser,
+    build_tab_completion_table
+)
 
 
 GLOBAL_ALIAS_PATH = os.path.join(GLOBAL_CONFIG_DIR, ALIAS_FILE_NAME)
@@ -121,8 +126,8 @@ def transform(self, args):
         # Only load the entire command table if it detects changes in the alias config
         if self.detect_alias_config_change():
             self.load_full_command_table()
-            self.collided_alias = AliasManager.build_collision_table(self.alias_table.sections(),
-                                                                     azext_alias.cached_reserved_commands)
+            self.collided_alias = AliasManager.build_collision_table(self.alias_table.sections())
+            build_tab_completion_table(self.alias_table)
         else:
             self.load_collided_alias()
 
@@ -220,7 +225,7 @@ def parse_error(self):
         return not self.alias_table.sections() and self.alias_config_str
 
     @staticmethod
-    def build_collision_table(aliases, reserved_commands, levels=COLLISION_CHECK_LEVEL_DEPTH):
+    def build_collision_table(aliases, levels=COLLISION_CHECK_LEVEL_DEPTH):
         """
         Build the collision table according to the alias configuration file against the entire command table.
 
@@ -246,7 +251,8 @@ def build_collision_table(aliases, reserved_commands, levels=COLLISION_CHECK_LEV
             word = alias.split()[0]
             for level in range(1, levels + 1):
                 collision_regex = r'^{}{}($|\s)'.format(r'([a-z\-]*\s)' * (level - 1), word.lower())
-                if list(filter(re.compile(collision_regex).match, reserved_commands)):
+                if list(filter(re.compile(collision_regex).match, azext_alias.cached_reserved_commands)) \
+                        and level not in collided_alias[word]:
                     collided_alias[word].append(level)
 
         telemetry.set_collided_aliases(list(collided_alias.keys()))

src/alias/azext_alias/azext_metadata.json

@@ -1,4 +1,4 @@
 {
-    "azext.minCliCoreVersion": "2.0.28",
+    "azext.minCliCoreVersion": "2.0.31.dev0",
     "azext.isPreview": true
 }

src/alias/azext_alias/custom.py

@@ -7,10 +7,9 @@
 
 from knack.util import CLIError
 
-import azext_alias
 from azext_alias._const import ALIAS_NOT_FOUND_ERROR
 from azext_alias.alias import GLOBAL_ALIAS_PATH, AliasManager
-from azext_alias.util import get_alias_table
+from azext_alias.util import get_alias_table, build_tab_completion_table
 
 
 def create_alias(alias_name, alias_command):
@@ -78,6 +77,6 @@ def _commit_change(alias_table):
         alias_config_hash = hashlib.sha1(alias_config_file.read().encode('utf-8')).hexdigest()
         AliasManager.write_alias_config_hash(alias_config_hash)
 
-    collided_alias = AliasManager.build_collision_table(alias_table.sections(),
-                                                        azext_alias.cached_reserved_commands)
+    collided_alias = AliasManager.build_collision_table(alias_table.sections())
     AliasManager.write_collided_alias(collided_alias)
+    build_tab_completion_table(alias_table)

src/alias/azext_alias/hooks.py

@@ -0,0 +1,143 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import json
+import timeit
+
+from knack.log import get_logger
+
+from azure.cli.command_modules.interactive.azclishell.command_tree import CommandBranch
+from azext_alias import telemetry
+from azext_alias.alias import AliasManager
+from azext_alias.util import (
+    is_alias_create_command,
+    cache_reserved_commands,
+    get_alias_table,
+    filter_aliases
+)
+from azext_alias._const import DEBUG_MSG_WITH_TIMING, GLOBAL_ALIAS_TAB_COMP_TABLE_PATH
+
+logger = get_logger(__name__)
+
+
+def alias_event_handler(_, **kwargs):
+    """
+    An event handler for alias transformation when EVENT_INVOKER_PRE_TRUNCATE_CMD_TBL event is invoked.
+    """
+    try:
+        telemetry.start()
+
+        start_time = timeit.default_timer()
+        args = kwargs.get('args')
+        alias_manager = AliasManager(**kwargs)
+
+        # [:] will keep the reference of the original args
+        args[:] = alias_manager.transform(args)
+
+        if is_alias_create_command(args):
+            load_cmd_tbl_func = kwargs.get('load_cmd_tbl_func', lambda _: {})
+            cache_reserved_commands(load_cmd_tbl_func)
+
+        elapsed_time = (timeit.default_timer() - start_time) * 1000
+        logger.debug(DEBUG_MSG_WITH_TIMING, args, elapsed_time)
+
+        telemetry.set_execution_time(round(elapsed_time, 2))
+    except Exception as client_exception:  # pylint: disable=broad-except
+        telemetry.set_exception(client_exception)
+        raise
+    finally:
+        telemetry.conclude()
+
+
+def enable_aliases_autocomplete(_, **kwargs):
+    """
+    Enable aliases autocomplete by injecting aliases into Azure CLI tab completion list.
+    """
+    external_completions = kwargs.get('external_completions', [])
+    prefix = kwargs.get('cword_prefix', [])
+    cur_commands = kwargs.get('comp_words', [])
+    alias_table = get_alias_table()
+    # Transform aliases if they are in current commands,
+    # so parser can get the correct subparser when chaining aliases
+    _transform_cur_commands(cur_commands, alias_table=alias_table)
+
+    for alias, alias_command in filter_aliases(alias_table):
+        if alias.startswith(prefix) and alias.strip() != prefix and _is_autocomplete_valid(cur_commands, alias_command):
+            # Only autocomplete the first word because alias is space-delimited
+            external_completions.append(alias)
+
+    # Append spaces if necessary (https://github.com/kislyuk/argcomplete/blob/master/argcomplete/__init__.py#L552-L559)
+    prequote = kwargs.get('cword_prequote', '')
+    continuation_chars = "=/:"
+    if len(external_completions) == 1 and external_completions[0][-1] not in continuation_chars and not prequote:
+        external_completions[0] += ' '
+
+
+def transform_cur_commands_interactive(_, **kwargs):
+    """
+    Transform any aliases in current commands in interactive into their respective commands.
+    """
+    event_payload = kwargs.get('event_payload', {})
+    # text_split = current commands typed in the interactive shell without any unfinished word
+    # text = current commands typed in the interactive shell
+    cur_commands = event_payload.get('text', '').split(' ')
+    _transform_cur_commands(cur_commands)
+
+    event_payload.update({
+        'text': ' '.join(cur_commands)
+    })
+
+
+def enable_aliases_autocomplete_interactive(_, **kwargs):
+    """
+    Enable aliases autocomplete on interactive mode by injecting aliases in the command tree.
+    """
+    subtree = kwargs.get('subtree', None)
+    if not subtree or not hasattr(subtree, 'children'):
+        return
+
+    for alias, alias_command in filter_aliases(get_alias_table()):
+        # Only autocomplete the first word because alias is space-delimited
+        if subtree.in_tree(alias_command.split()):
+            subtree.add_child(CommandBranch(alias))
+
+
+def _is_autocomplete_valid(cur_commands, alias_command):
+    """
+    Determine whether autocomplete can be performed at the current state.
+
+    Args:
+        parser: The current CLI parser.
+        cur_commands: The current commands typed in the console.
+        alias_command: The alias command.
+
+    Returns:
+        True if autocomplete can be performed.
+    """
+    parent_command = ' '.join(cur_commands[1:])
+    with open(GLOBAL_ALIAS_TAB_COMP_TABLE_PATH, 'r') as tab_completion_table_file:
+        try:
+            tab_completion_table = json.loads(tab_completion_table_file.read())
+            return alias_command in tab_completion_table and parent_command in tab_completion_table[alias_command]
+        except Exception:  # pylint: disable=broad-except
+            return False
+
+
+def _transform_cur_commands(cur_commands, alias_table=None):
+    """
+    Transform any aliases in cur_commands into their respective commands.
+
+    Args:
+        alias_table: The alias table.
+        cur_commands: current commands typed in the console.
+    """
+    transformed = []
+    alias_table = alias_table if alias_table else get_alias_table()
+    for cmd in cur_commands:
+        if cmd in alias_table.sections() and alias_table.has_option(cmd, 'command'):
+            transformed += alias_table.get(cmd, 'command').split()
+        else:
+            transformed.append(cmd)
+    cur_commands[:] = transformed