Complete Completion (#88)

Author: Juliya Smith

CHANGELOG.md

@@ -84,6 +84,8 @@ how a consumer would use the library (e.g. adding unit tests, updating documenta
 
 - Short option `-u` added for `code42 high-risk-employee add-risk-tags` and `remove-risk-tags`.
 
+- Tab completion for bash and zsh for Unix based machines.
+
 ### Fixed
 
 - Fixed bug in bulk commands where value-less fields in csv files were treated as empty strings instead of None.

src/code42cli/cmds/alerts/rules/commands.py

@@ -131,7 +131,7 @@ def load_commands(self):
         add = Command(
             self.ADD_USER,
             u"Add a user to an alert rule.",
-            u"{} add-user --rule-id <id>  --username <username>".format(usage_prefix),
+            u"{} add-user --rule-id <id> --username <username>".format(usage_prefix),
             handler=add_user,
             arg_customizer=_customize_add_arguments,
         )

src/code42cli/cmds/detectionlists/__init__.py

@@ -144,16 +144,16 @@ def generate_template_file(self, cmd, path=None):
         handler = detection_list.get_handler(self.handlers, cmd)
         generate_template(handler, path)
 
-    def bulk_add_employees(self, sdk, profile, csv_file):
+    def bulk_add_employees(self, sdk, profile, filename):
         """Takes a csv file with each row representing an employee and adds them all to a 
         detection list in a bulk fashion.
 
         Args:
             sdk (py42.sdk.SDKClient): The py42 sdk.
             profile (Code42Profile): The profile under which to execute this command.
-            csv_file (str or unicode): The path to the csv file containing rows of users.
+            filename (str or unicode): The path to the csv file containing rows of users.
         """
-        reader = create_csv_reader(csv_file)
+        reader = create_csv_reader(filename)
         run_bulk_process(lambda **kwargs: self._add_employee(sdk, profile, **kwargs), reader)
 
     def bulk_remove_employees(self, sdk, profile, users_file):
@@ -176,12 +176,12 @@ def _add_employee(self, sdk, profile, **kwargs):
     def _remove_employee(self, sdk, profile, *args, **kwargs):
         self.handlers.remove_employee(sdk, profile, *args, **kwargs)
 
-    def bulk_add_risk_tags(self, sdk, profile, csv_file):
-        reader = create_csv_reader(csv_file)
+    def bulk_add_risk_tags(self, sdk, profile, filename):
+        reader = create_csv_reader(filename)
         run_bulk_process(lambda **kwargs: add_risk_tags(sdk, profile, **kwargs), reader)
 
-    def bulk_remove_risk_tags(self, sdk, profile, csv_file):
-        reader = create_csv_reader(csv_file)
+    def bulk_remove_risk_tags(self, sdk, profile, filename):
+        reader = create_csv_reader(filename)
         run_bulk_process(lambda **kwargs: remove_risk_tags(sdk, profile, **kwargs), reader)
 
 

src/code42cli/cmds/detectionlists/commands.py

@@ -146,8 +146,8 @@ def _load_hre_bulk_generate_template_description(argument_collection):
         cmd_type.set_choices(HighRiskBulkCommandType())
 
     def _load_bulk_add_description(self, argument_collection):
-        csv_file = argument_collection.arg_configs[u"csv_file"]
-        csv_file.set_help(
+        filename = argument_collection.arg_configs[u"filename"]
+        filename.set_help(
             u"The path to the csv file for bulk adding users to the {} detection list.".format(
                 self._name
             )
@@ -162,16 +162,16 @@ def _load_bulk_remove_description(self, argument_collection):
         )
 
     def _load_bulk_add_risk_tags_description(self, argument_collection):
-        csv_file = argument_collection.arg_configs[u"csv_file"]
-        csv_file.set_help(
+        filename = argument_collection.arg_configs[u"filename"]
+        filename.set_help(
             u"A file containing a ',' separated username with space-separated tags to add "
             u"to the {} detection list. "
             u"e.g. [email protected],tag1 tag2 tag3".format(self._name)
         )
 
     def _load_bulk_remove_risk_tags_description(self, argument_collection):
-        csv_file = argument_collection.arg_configs[u"csv_file"]
-        csv_file.set_help(
+        filename = argument_collection.arg_configs[u"filename"]
+        filename.set_help(
             u"A file containing a ',' separated username with space-separated tags to remove "
             u"from the {} detection list. "
             u"e.g. [email protected],tag1 tag2 tag3".format(self._name)

src/code42cli/cmds/legal_hold/__init__.py

@@ -59,16 +59,12 @@ def get_matters(sdk):
 
 def add_bulk_users(sdk, file_name):
     reader = create_csv_reader(file_name)
-    run_bulk_process(
-        lambda matter_id, username: add_user(sdk, matter_id, username), reader,
-    )
+    run_bulk_process(lambda matter_id, username: add_user(sdk, matter_id, username), reader)
 
 
 def remove_bulk_users(sdk, file_name):
     reader = create_csv_reader(file_name)
-    run_bulk_process(
-        lambda matter_id, username: remove_user(sdk, matter_id, username), reader,
-    )
+    run_bulk_process(lambda matter_id, username: remove_user(sdk, matter_id, username), reader)
 
 
 def show_matter(sdk, matter_id, include_inactive=False, include_policy=False):

src/code42cli/cmds/securitydata/main.py

@@ -100,7 +100,6 @@ def _load_send_to_args(arg_collection):
             help=u"Protocol used to send logs to server.",
         ),
     }
-
     arg_collection.extend(send_to_args)
     _load_search_args(arg_collection)
 

src/code42cli/commands.py

@@ -3,6 +3,7 @@
 from code42cli import profile as cliprofile
 from code42cli.args import get_auto_arg_configs, SDK_ARG_NAME, PROFILE_ARG_NAME
 from code42cli.sdk_client import create_sdk
+from code42cli.tree_nodes import SubcommandNode
 
 
 class DictObject(object):
@@ -150,29 +151,24 @@ def _kvps_to_obj(kvps):
 
 
 class SubcommandLoader(object):
-    """Responsible for creating subcommands for it's root command. It is also useful for getting 
-    command information ahead of time, as in the example of tab completion."""
+    """Responsible for creating subcommands for it's root command."""
 
-    def __init__(self, root_command_name):
+    def __init__(self, root_command_name, node=None):
         self.root = root_command_name
+        self._node = node
 
-    @property
-    def names(self):
-        """The names of all the subcommands in this subcommabd loader's root command."""
-        sub_cmds = self.load_commands()
-        return [cmd.name for cmd in sub_cmds]
+    def __getitem__(self, item):
+        return self.get_node()[item]
 
     @property
-    def subtrees(self):
-        """All subcommands for this subcommand loader's root command mapped to their given 
-        subcommand loaders."""
-        cmds = self.load_commands()
-        results = {}
-        for cmd in cmds:
-            subcommand_loader = cmd.subcommand_loader
-            if subcommand_loader:
-                results[cmd.name] = subcommand_loader
-        return results
+    def names(self):
+        return self.get_node().names
 
     def load_commands(self):
+        """Override"""
         return []
+
+    def get_node(self):
+        if not self._node:
+            self._node = SubcommandNode(self.root, self.load_commands())
+        return self._node

src/code42cli/completer.py

@@ -1,5 +1,9 @@
+from os import path
+
 from code42cli import MAIN_COMMAND
 from code42cli.main import MainSubcommandLoader
+from code42cli.tree_nodes import ArgNode
+from code42cli.util import get_files_in_path
 
 
 def _get_matches(current, options):
@@ -11,46 +15,67 @@ def _get_matches(current, options):
     return matches
 
 
-def _get_next_full_set_of_commands(cmd_loader, current):
-    cmd_loader = cmd_loader.subtrees[current]
-    return cmd_loader.names
+def _get_next_full_set_of_options(node, current):
+    node = node[current]
+    names = list(node.names)
+    if _can_complete_with_local_files(current, node):
+        files = get_files_in_path("")
+        names.extend(files)
+    return names
+
+
+def _can_complete_with_local_files(current, node):
+    return isinstance(node, ArgNode) and (not current or current[0] != u"-")
 
 
 class Completer(object):
     def __init__(self, main_cmd_loader=None):
-        self._main_cmd_loader = main_cmd_loader or MainSubcommandLoader(u"")
+        self._main_cmd_loader = main_cmd_loader or MainSubcommandLoader()
 
     def complete(self, cmdline, point=None):
         try:
             point = point or len(cmdline)
             args = cmdline[0:point].split()
+            # Complete with main commands if `code42` is typed out.
+            # Note that the command `code42` should complete on its own.
             if len(args) < 2:
-                # `code42` already completes w/o
                 return self._main_cmd_loader.names if args[0] == MAIN_COMMAND else []
 
             current = args[-1]
-            cmd_loader = self._search_trees(args)
-            if not cmd_loader:
-                return []
+            search_results, options = self._get_completion_options(args)
 
-            options = cmd_loader.names
+            # Complete with full set of arg/command options
             if current in options:
-                # `current` is already complete
-                return _get_next_full_set_of_commands(cmd_loader, current)
+                return _get_next_full_set_of_options(search_results, current)
+
+            if _can_complete_with_local_files(current, search_results):
+                files = get_files_in_path(current)
+                if current[0] == "~":
+                    replace = path.expanduser("~")
+                    files = [f.replace(replace, "~") for f in files]
+                options.extend(files)
 
             return _get_matches(current, options) if options else []
         except:
             return []
 
     def _search_trees(self, args):
         # Find cmd_loader at lowest level from given args
-        cmd_loader = self._main_cmd_loader
+        node = self._main_cmd_loader.get_node()
         if len(args) > 2:
             for arg in args[1:-1]:
-                cmd_loader = cmd_loader.subtrees[arg]
-        return cmd_loader
+                next_node = node[arg]
+                if next_node:
+                    node = next_node
+                else:
+                    return node
+        return node
+
+    def _get_completion_options(self, args):
+        search_results = self._search_trees(args)
+        return search_results, search_results.names
 
 
 def complete(cmdline, point):
-    choices = Completer().complete(cmdline, point)
+    choices = Completer().complete(cmdline, point) or []
     print(u" \n".join(choices))

src/code42cli/main.py

@@ -53,6 +53,9 @@ class MainSubcommandLoader(SubcommandLoader):
     HIGH_RISK_EMPLOYEE = DetectionLists.HIGH_RISK_EMPLOYEE
     LEGAL_HOLD = u"legal-hold"
 
+    def __init__(self):
+        super(MainSubcommandLoader, self).__init__(u"")
+
     def load_commands(self):
         detection_lists_description = (
             u"For adding and removing employees from the {} detection list."
@@ -118,7 +121,7 @@ def _create_legal_hold_loader(self):
 
 
 def main():
-    top = Command(u"", u"", subcommand_loader=MainSubcommandLoader(u""))
+    top = Command(u"", u"", subcommand_loader=MainSubcommandLoader())
     invoker = CommandInvoker(top)
     try:
         invoker.run(sys.argv[1:])

src/code42cli/tree_nodes.py

@@ -0,0 +1,104 @@
+class CLINode(object):
+    """Base class for identifying nodes in the command/argument hierarchy."""
+
+    @property
+    def names(self):
+        """Override"""
+        return []
+
+
+class ChoicesNode(CLINode):
+    """A node who `names` refer to choices the user can select for an argument."""
+
+    def __init__(self, options):
+        self._choices = options
+
+    def __iter__(self):
+        return iter(self._choices)
+
+    def __getitem__(self, item):
+        return self._choices[item]
+
+    def get(self, item):
+        return self._choices.get(item)
+
+    @property
+    def names(self):
+        return self._choices
+
+
+class ArgNode(CLINode):
+    """A node whose `names` are a list of flagged arguments the user can select from."""
+
+    def __init__(self, args):
+        self.args = args
+
+    @property
+    def names(self):
+        try:
+            arg_names = [
+                n
+                for names in [self.args[key].settings[u"options_list"] for key in self.args]
+                for n in names
+                if n.startswith("--")
+            ]
+            return arg_names
+        except:
+            return self.args
+
+    def __getitem__(self, item):
+        """Access sub loaders to navigate the argument/options tree, connected to a leaf command."""
+        if item in self.args:
+            return ArgNode(self.args)
+
+        for key in self.args:
+            arg = self.args[key]
+            if item not in arg.settings[u"options_list"]:
+                continue
+            choices = arg.settings[u"choices"]
+            if choices:
+                return ChoicesNode(choices)
+        return ArgNode(self.args)
+
+    def __iter__(self):
+        return iter(self.names)
+
+
+class SubcommandNode(CLINode):
+    """Gets command information ahead of command-execution."""
+
+    def __init__(self, root_command_name, commands):
+        self.root = root_command_name
+        self.commands = commands
+
+    def __getitem__(self, item):
+        try:
+            return self._subtrees[item]
+        except KeyError:
+            return self._get_args(item)
+
+    def _get_args(self, item):
+        cmd = self._get_command_by_name(item)
+        if cmd:
+            args = cmd.get_arg_configs()
+            return ArgNode(args)
+
+    def _get_command_by_name(self, name):
+        for cmd in self.commands:
+            if cmd.name == name:
+                return cmd
+
+    @property
+    def names(self):
+        """The names of all the subcommands in this subcommand loader's root command."""
+        return [cmd.name for cmd in self.commands]
+
+    @property
+    def _subtrees(self):
+        """Maps subcommand names to their respective subcommand nodes."""
+        results = {}
+        for cmd in self.commands:
+            if cmd.subcommand_loader:
+                commands = cmd.subcommand_loader.load_commands()
+                results[cmd.name] = SubcommandNode(cmd.name, commands)
+        return results