Add support for custom aliases for commands (#755)

Author: ezilber-akamai

linodecli/__init__.py

@@ -46,10 +46,7 @@
 
 # if any of these arguments are given, we don't need to prompt for configuration
 skip_config = (
-    any(
-        c in argv
-        for c in ["--skip-config", "--help", "--version", "completion"]
-    )
+    any(c in argv for c in ["--skip-config", "--version", "completion"])
     or TEST_MODE
 )
 
@@ -110,6 +107,43 @@ def main():  # pylint: disable=too-many-branches,too-many-statements
         # if not spec was found and we weren't baking, we're doomed
         sys.exit(ExitCodes.ARGUMENT_ERROR)
 
+    if parsed.command in ("set-custom-alias", "remove-custom-alias"):
+        if not parsed.target_command or not parsed.alias:
+            print(
+                "Both --target-command and --alias must be provided.",
+                file=sys.stderr,
+            )
+            sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        command = parsed.target_command
+        alias = parsed.alias
+
+        if command not in cli.ops:
+            print(
+                f"Error: '{command}' is not a valid command.", file=sys.stderr
+            )
+            sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        if parsed.command == "set-custom-alias":
+            if (alias, command) not in cli.config.get_custom_aliases().items():
+                cli.config.set_custom_alias(alias, command)
+                print(f"Custom alias '{alias}' set for command '{command}'")
+            else:
+                print(
+                    f"Custom alias '{alias}' already set for command '{command}'"
+                )
+
+        if parsed.command == "remove-custom-alias":
+            if (alias, command) in cli.config.get_custom_aliases().items():
+                cli.config.remove_custom_alias(alias, command)
+                print(f"Custom alias '{alias}' removed for command '{command}'")
+            else:
+                print(
+                    f"Custom alias '{alias}' does not exist for command '{command}'"
+                )
+
+        sys.exit(ExitCodes.SUCCESS)
+
     if parsed.command == "register-plugin":
         if parsed.action is None:
             print("register-plugin requires a module name!", file=sys.stderr)
@@ -213,28 +247,36 @@ def main():  # pylint: disable=too-many-branches,too-many-statements
         plugin_args.remove(parsed.command)  # don't include the plugin name
         plugins.invoke(parsed.command, plugin_args, context)
         sys.exit(ExitCodes.SUCCESS)
-
     # unknown commands
     if (
         parsed.command not in cli.ops
         and parsed.command not in plugins.available(cli.config)
         and parsed.command not in HELP_TOPICS
+        and parsed.command not in cli.config.get_custom_aliases().keys()
     ):
         print(f"Unrecognized command {parsed.command}", file=sys.stderr)
         sys.exit(ExitCodes.UNRECOGNIZED_COMMAND)
 
     # handle a help for a command - either --help or no action triggers this
-    if (
-        parsed.command is not None
-        and parsed.action is None
-        and parsed.command in cli.ops
-    ):
-        print_help_command_actions(cli.ops, parsed.command)
-        sys.exit(ExitCodes.SUCCESS)
+    if parsed.command is not None and parsed.action is None:
+        if parsed.command in cli.ops:
+            print_help_command_actions(cli.ops, parsed.command)
+            sys.exit(ExitCodes.SUCCESS)
+        if parsed.command in cli.config.get_custom_aliases().keys():
+            print_help_command_actions(
+                cli.ops, cli.config.get_custom_aliases()[parsed.command]
+            )
 
     if parsed.command is not None and parsed.action is not None:
         if parsed.help:
-            print_help_action(cli, parsed.command, parsed.action)
+            if parsed.command in cli.config.get_custom_aliases().keys():
+                print_help_action(
+                    cli,
+                    cli.config.get_custom_aliases()[parsed.command],
+                    parsed.action,
+                )
+            else:
+                print_help_action(cli, parsed.command, parsed.action)
             sys.exit(ExitCodes.SUCCESS)
 
         cli.handle_command(parsed.command, parsed.action, args)

linodecli/arg_helpers.py

@@ -67,6 +67,22 @@ def register_args(parser: ArgumentParser) -> ArgumentParser:
         help="Prints version information and exits.",
     )
 
+    parser.add_argument(
+        "--target-command",
+        "-t",
+        nargs="?",
+        type=str,
+        help="The command to set or remove an alias for.",
+    )
+
+    parser.add_argument(
+        "--alias",
+        "-a",
+        nargs="?",
+        type=str,
+        help="The alias to set or remove.",
+    )
+
     # Register shared argument groups
     register_output_args_shared(parser)
     register_pagination_args_shared(parser)

linodecli/cli.py

@@ -259,7 +259,11 @@ def find_operation(self, command, action):
         Finds the corresponding operation for the given command and action.
         """
         if command not in self.ops:
-            raise ValueError(f"Command not found: {command}")
+            # Check that the passed command is not an alias before raising an error
+            if command in self.config.get_custom_aliases().keys():
+                command = self.config.get_custom_aliases()[command]
+            else:
+                raise ValueError(f"Command not found: {command}")
 
         command_dict = self.ops[command]
 

linodecli/configuration/config.py

@@ -597,3 +597,60 @@ def _handle_no_default_user(self):  # pylint: disable=too-many-branches
                 self.write_config()
                 return
             print(f"No user {username}")
+
+    def set_custom_alias(self, alias: str, command: str):
+        """
+        Sets a custom alias for a Linode CLI command.
+
+        :param alias: The custom alias name.
+        :param command: The command the custom alias maps to.
+        """
+        if not self.config.has_section("custom_aliases"):
+            self.config.add_section("custom_aliases")
+
+        self.config.set("custom_aliases", alias, command)
+        self.write_config()
+
+    def remove_custom_alias(self, alias: str, command: str):
+        """
+        Removes a custom alias from the Linode CLI configuration.
+
+        :param alias: The alias name to remove.
+        :param command: The command the alias is mapped to.
+        """
+        if not self.config.has_section("custom_aliases"):
+            print("Error: No custom aliases have been set.", file=sys.stderr)
+            return
+
+        if not self.config.has_option("custom_aliases", alias):
+            print(
+                f"Error: Custom alias '{alias}' does not exist.",
+                file=sys.stderr,
+            )
+            return
+
+        # Check if the alias maps to the given command
+        existing_command = self.config.get("custom_aliases", alias)
+        if existing_command != command:
+            print(
+                f"Error: Custom alias '{alias}' is mapped to "
+                f"'{existing_command}', not '{command}'.",
+                file=sys.stderr,
+            )
+            return
+
+        # Remove the alias and update the config
+        self.config.remove_option("custom_aliases", alias)
+        self.write_config()
+
+    def get_custom_aliases(self) -> Dict[str, str]:
+        """
+        Retrieves all stored custom command aliases from the config.
+
+        :return: A dictionary mapping custom alias names to their respective commands.
+        """
+        return (
+            dict(self.config.items("custom_aliases"))
+            if (self.config.has_section("custom_aliases"))
+            else {}
+        )