Polishing argcomplete support

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()