{Core} Improve autocomplete performance by preventing loading all modules (#24765)

Author: bebound

src/azure-cli-core/azure/cli/core/__init__.py

@@ -449,6 +449,20 @@ def _get_extension_suppressions(mod_loaders):
                 if raw_cmd in self.command_group_table:
                     logger.debug("Found a match in the command group table for '%s'.", raw_cmd)
                     return self.command_table
+                if self.cli_ctx.data['completer_active']:
+                    # If the command is not complete in autocomplete mode, we should match shorter command.
+                    # For example, `account sho` should match `account`.
+                    logger.debug("Could not find a match in the command or command group table for '%s'", raw_cmd)
+                    trimmed_raw_cmd = ' '.join(raw_cmd.split()[:-1])
+                    logger.debug("In autocomplete mode, try to match trimmed raw cmd: '%s'", trimmed_raw_cmd)
+                    if not trimmed_raw_cmd:
+                        # If full command is 'az acc', raw_cmd is 'acc', trimmed_raw_cmd is ''.
+                        logger.debug("Trimmed raw cmd is empty, return command table.")
+                        return self.command_table
+                    if trimmed_raw_cmd in self.command_group_table:
+                        logger.debug("Found a match in the command group table for trimmed raw cmd: '%s'.",
+                                     trimmed_raw_cmd)
+                        return self.command_table
 
                 logger.debug("Could not find a match in the command or command group table for '%s'. "
                              "The index may be outdated.", raw_cmd)
@@ -527,6 +541,7 @@ def __init__(self, cli_ctx=None):
         if cli_ctx:
             self.version = __version__
             self.cloud_profile = cli_ctx.cloud.profile
+        self.cli_ctx = cli_ctx
 
     def get(self, args):
         """Get the corresponding module and extension list of a command.
@@ -555,6 +570,13 @@ def get(self, args):
         # Check the command index for (command: [module]) mapping, like
         # "network": ["azure.cli.command_modules.natgateway", "azure.cli.command_modules.network", "azext_firewall"]
         index_modules_extensions = index.get(top_command)
+        if not index_modules_extensions and self.cli_ctx.data['completer_active']:
+            # If user type `az acco`, command begin with `acco` will be matched.
+            logger.debug("In autocomplete mode, load commands starting with: '%s'", top_command)
+            index_modules_extensions = []
+            for command in index:
+                if command.startswith(top_command):
+                    index_modules_extensions += index[command]
 
         if index_modules_extensions:
             # This list contains both built-in modules and extensions