Polishing argcomplete support (#829)

* Polishing argcomplete support

* Forgot to add license notice on new files
* Add support for completing subcommands. Its not easy to
  do this with argparse's subparser with the way traitlets
  uses argparse, so instead directly add subcommands as
  completions for the first cword in argcomplete.
* A few other small refactor/cleanup/testing

* Add missed subapp.clear_instance() in unit test

Otherwise, this test prevents other tests from creating
subcommand applications.

Author: azjps

traitlets/config/application.py

@@ -731,7 +731,7 @@ def flatten_flags(self):
         This prevents issues such as an alias pointing to InteractiveShell,
         but a config file setting the same trait in TerminalInteraciveShell
         getting inappropriate priority over the command-line arg.
-        Also, loaders expect ``(key: longname)`` and not ````key: (longname, help)`` items.
+        Also, loaders expect ``(key: longname)`` and not ``key: (longname, help)`` items.
 
         Only aliases with exactly one descendent in the class list
         will be promoted.
@@ -785,7 +785,9 @@ def flatten_flags(self):
         return flags, aliases
 
     def _create_loader(self, argv, aliases, flags, classes):
-        return KVArgParseConfigLoader(argv, aliases, flags, classes=classes, log=self.log)
+        return KVArgParseConfigLoader(
+            argv, aliases, flags, classes=classes, log=self.log, subcommands=self.subcommands
+        )
 
     @classmethod
     def _get_sys_argv(cls, check_argcomplete: bool = False) -> t.List[str]:

traitlets/config/argcomplete_config.py

@@ -1,4 +1,9 @@
 """Helper utilities for integrating argcomplete with traitlets"""
+
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+
 import argparse
 import os
 import typing as t
@@ -76,23 +81,25 @@ def increment_argcomplete_index():
 class ExtendedCompletionFinder(CompletionFinder):
     """An extension of CompletionFinder which dynamically completes class-trait based options
 
-    This finder mainly adds 2 functionalities:
+    This finder adds a few functionalities:
 
-    1. When completing options, it will add --Class. to the list of completions, for each
-    class in Application.classes that could complete the current option.
-    2. If it detects that we are currently trying to complete an option related to --Class.,
-    it will add the corresponding config traits of Class to the ArgumentParser instance,
+    1. When completing options, it will add ``--Class.`` to the list of completions, for each
+    class in `Application.classes` that could complete the current option.
+    2. If it detects that we are currently trying to complete an option related to ``--Class.``,
+    it will add the corresponding config traits of Class to the `ArgumentParser` instance,
     so that the traits' completers can be used.
+    3. If there are any subcommands, they are added as completions for the first word
 
-    Note that we are avoiding adding all config traits of all classes to the ArgumentParser,
+    Note that we are avoiding adding all config traits of all classes to the `ArgumentParser`,
     which would be easier but would add more runtime overhead and would also make completions
     appear more spammy.
 
-    These changes do require using the internals of argcomplete.CompletionFinder.
+    These changes do require using the internals of `argcomplete.CompletionFinder`.
     """
 
     _parser: argparse.ArgumentParser
-    config_classes: t.List[t.Any]  # Configurables
+    config_classes: t.List[t.Any] = []  # Configurables
+    subcommands: t.List[str] = []
 
     def match_class_completions(self, cword_prefix: str) -> t.List[t.Tuple[t.Any, str]]:
         """Match the word to be completed against our Configurable classes
@@ -182,6 +189,16 @@ def _get_completions(
 
         completions: t.List[str]
         completions = super()._get_completions(comp_words, cword_prefix, *args)
+
+        # For subcommand-handling: it is difficult to get this to work
+        # using argparse subparsers, because the ArgumentParser accepts
+        # arbitrary extra_args, which ends up masking subparsers.
+        # Instead, check if comp_words only consists of the script,
+        # if so check if any subcommands start with cword_prefix.
+        if self.subcommands and len(comp_words) == 1:
+            argcomplete.debug("Adding subcommands for", cword_prefix)
+            completions.extend(subc for subc in self.subcommands if subc.startswith(cword_prefix))
+
         return completions
 
     def _get_option_completions(

traitlets/config/loader.py

@@ -784,11 +784,15 @@ def parse_known_args(self, args=None, namespace=None):
         return super().parse_known_args(args, namespace)
 
 
+# type aliases
+Flags = t.Union[str, t.Tuple[str, ...]]
+SubcommandsDict = t.Dict[str, t.Any]
+
+
 class ArgParseConfigLoader(CommandLineConfigLoader):
     """A loader that uses the argparse module to load from the command line."""
 
     parser_class = ArgumentParser
-    Flags = t.Union[str, t.Tuple[str, ...]]
 
     def __init__(
         self,
@@ -797,6 +801,7 @@ def __init__(
         flags: t.Optional[t.Dict[Flags, str]] = None,
         log: t.Any = None,
         classes: t.Optional[t.List[t.Type[t.Any]]] = None,
+        subcommands: t.Optional[SubcommandsDict] = None,
         *parser_args: t.Any,
         **parser_kw: t.Any,
     ) -> None:
@@ -837,6 +842,7 @@ def __init__(
         self.aliases = aliases or {}
         self.flags = flags or {}
         self.classes = classes
+        self.subcommands = subcommands  # only used for argcomplete currently
 
         self.parser_args = parser_args
         self.version = parser_kw.pop("version", None)
@@ -874,6 +880,7 @@ def load_config(self, argv=None, aliases=None, flags=_deprecated, classes=None):
         if classes is not None:
             self.classes = classes
         self._create_parser()
+        self._argcomplete(self.classes, self.subcommands)
         self._parse_args(argv)
         self._convert_to_config()
         return self.config
@@ -893,6 +900,12 @@ def _create_parser(self):
     def _add_arguments(self, aliases, flags, classes):
         raise NotImplementedError("subclasses must implement _add_arguments")
 
+    def _argcomplete(
+        self, classes: t.List[t.Any], subcommands: t.Optional[SubcommandsDict]
+    ) -> None:
+        """If argcomplete is enabled, allow triggering command-line autocompletion"""
+        pass
+
     def _parse_args(self, args):
         """self.parser->self.parsed_data"""
         uargs = [cast_unicode(a) for a in args]
@@ -1047,7 +1060,6 @@ def _add_arguments(self, aliases, flags, classes):
                 if argcompleter is not None:
                     # argcomplete's completers are callables returning list of completion strings
                     action.completer = functools.partial(argcompleter, key=key)  # type: ignore
-        self.argcomplete(classes)
 
     def _convert_to_config(self):
         """self.parsed_data->self.config, parse unrecognized extra args via KVLoader."""
@@ -1097,7 +1109,10 @@ def _handle_unrecognized_alias(self, arg: str) -> None:
         """
         self.log.warning("Unrecognized alias: '%s', it will have no effect.", arg)
 
-    def argcomplete(self, classes: t.List[t.Any]) -> None:
+    def _argcomplete(
+        self, classes: t.List[t.Any], subcommands: t.Optional[SubcommandsDict]
+    ) -> None:
+        """If argcomplete is enabled, allow triggering command-line autocompletion"""
         try:
             import argcomplete  # type: ignore[import]  # noqa
         except ImportError:
@@ -1107,6 +1122,7 @@ def argcomplete(self, classes: t.List[t.Any]) -> None:
 
         finder = argcomplete_config.ExtendedCompletionFinder()
         finder.config_classes = classes
+        finder.subcommands = list(subcommands or [])
         # for ease of testing, pass through self._argcomplete_kwargs if set
         finder(self.parser, **getattr(self, "_argcomplete_kwargs", {}))
 

traitlets/config/tests/test_application.py

@@ -633,6 +633,9 @@ def test_subcommands_instanciation(self):
         self.assertIs(app.subapp.parent, app)
         self.assertIs(app.subapp.subapp.parent, app.subapp)  # Set by factory.
 
+        Root.clear_instance()
+        Sub1.clear_instance()
+
     def test_loaded_config_files(self):
         app = MyApp()
         app.log = logging.getLogger()

traitlets/config/tests/test_argcomplete.py

@@ -2,6 +2,9 @@
 Tests for argcomplete handling by traitlets.config.application.Application
 """
 
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
+
 import io
 import os
 import typing as t
@@ -21,8 +24,19 @@ class ArgcompleteApp(Application):
 
     argcomplete_kwargs: t.Dict[str, t.Any]
 
+    def __init__(self, *args, **kwargs):
+        # For subcommands, inherit argcomplete_kwargs from parent app
+        parent = kwargs.get("parent")
+        super().__init__(*args, **kwargs)
+        if parent:
+            argcomplete_kwargs = getattr(parent, "argcomplete_kwargs", None)
+            if argcomplete_kwargs:
+                self.argcomplete_kwargs = argcomplete_kwargs
+
     def _create_loader(self, argv, aliases, flags, classes):
-        loader = KVArgParseConfigLoader(argv, aliases, flags, classes=classes, log=self.log)
+        loader = KVArgParseConfigLoader(
+            argv, aliases, flags, classes=classes, log=self.log, subcommands=self.subcommands
+        )
         loader._argcomplete_kwargs = self.argcomplete_kwargs  # type: ignore[attr-defined]
         return loader
 
@@ -169,24 +183,34 @@ class CustomApp(ArgcompleteApp):
         assert completions == ["--val=foo", "--val=bar"] or completions == ["foo", "bar"]
         assert self.run_completer(app, "app --val  --log-level=", point=10) == ["foo", "bar"]
 
-    # TODO: don't have easy way of testing subcommands yet, since we want
-    # to inject _argcomplete_kwargs to subapp. Could use mocking for this
-    # def test_complete_subcommands_subapp1(self, argcomplete_on):
-    #     # subcommand handling modifies _ARGCOMPLETE env var global state, so
-    #     # only can test one completion per unit test
-    #     app = MainApp()
-    #     assert set(self.run_completer(app, "app subapp1 --Sub")) > {
-    #         '--SubApp1.show_config',
-    #         '--SubApp1.log_level',
-    #         '--SubApp1.log_format',
-    #     }
-    #
-    # def test_complete_subcommands_subapp2(self, argcomplete_on):
-    #     app = MainApp()
-    #     assert set(self.run_completer(app, "app subapp2 --")) > {
-    #         '--Application.',
-    #         '--SubApp2.',
-    #     }
+    def test_complete_subcommands(self, argcomplete_on):
+        app = MainApp()
+        assert set(self.run_completer(app, "app ")) >= {"subapp1", "subapp2"}
+        assert set(self.run_completer(app, "app sub")) == {"subapp1", "subapp2"}
+        assert set(self.run_completer(app, "app subapp1")) == {"subapp1"}
+
+    def test_complete_subcommands_subapp1(self, argcomplete_on):
+        # subcommand handling modifies _ARGCOMPLETE env var global state, so
+        # only can test one completion per unit test
+        app = MainApp()
+        try:
+            assert set(self.run_completer(app, "app subapp1 --Sub")) > {
+                '--SubApp1.show_config',
+                '--SubApp1.log_level',
+                '--SubApp1.log_format',
+            }
+        finally:
+            SubApp1.clear_instance()
+
+    def test_complete_subcommands_subapp2(self, argcomplete_on):
+        app = MainApp()
+        try:
+            assert set(self.run_completer(app, "app subapp2 --")) > {
+                '--Application.',
+                '--SubApp2.',
+            }
+        finally:
+            SubApp2.clear_instance()
 
     def test_complete_subcommands_main(self, argcomplete_on):
         app = MainApp()