Added set-alias and remove-alias commands

Author: ezilber-akamai

linodecli/__init__.py

@@ -110,6 +110,53 @@ 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 == "set-custom-alias" or parsed.command == "remove-custom-alias":
+        # Ensure both --command and --alias are provided
+        # TODO: Can we use argparse for these
+        if len(args) == 4:
+            if args[0] != "--command" or args[2] != "--alias":
+                print("Usage: linode-cli set-custom-alias --command [COMMAND] --alias [ALIAS]", file=sys.stderr)
+                sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        # Get the indexes for --command and --alias
+        try:
+            command_index = args.index("--command") + 1
+            alias_index = args.index("--alias") + 1
+        except ValueError:
+            print("Both --command and --alias arguments are required.", file=sys.stderr)
+            sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        if command_index >= len(args) or alias_index >= len(args):
+            print("Both --command and --alias arguments must have valid values.", file=sys.stderr)
+            sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        # Retrieve the command and alias from arguments
+        command = args[command_index]
+        alias = args[alias_index]
+
+        # Check if the command is valid
+        if command not in cli.ops:
+            print(f"Error: '{command}' is not a valid command.", file=sys.stderr)
+            sys.exit(ExitCodes.ARGUMENT_ERROR)
+
+        # Set the alias if it does not already exist
+        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}'")
+
+        # Remove the alias if it already exists
+        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)
@@ -219,6 +266,7 @@ def main():  # pylint: disable=too-many-branches,too-many-statements
         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)

linodecli/cli.py

@@ -238,7 +238,13 @@ 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():
+                print(command)
+                command = self.config.get_custom_aliases()[command]
+                print(command)
+            else:
+                raise ValueError(f"Command not found: {command}")
 
         command_dict = self.ops[command]
 

linodecli/configuration/config.py

@@ -597,3 +597,49 @@ 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(f"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 '{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 {}