diff --git a/traitlets/config/argcomplete_config.py b/traitlets/config/argcomplete_config.py new file mode 100644 index 00000000..ebb56d38 --- /dev/null +++ b/traitlets/config/argcomplete_config.py @@ -0,0 +1,105 @@ +import typing as t + +import argcomplete + +class ExtendedCompletionFinder(argcomplete.CompletionFinder): + """An extension of CompletionFinder which dynamically completes class-trait based options + + This finder mainly adds 2 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, + so that the traits' completers can be used. + + 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. + """ + 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 + + Check if cword_prefix could potentially match against --{class}. for any class + in Application.classes. + """ + class_completions = [(cls, f"--{cls.__name__}.") for cls in self.config_classes] + matched_completions = class_completions + if "." in cword_prefix: + cword_prefix = cword_prefix[:cword_prefix.index(".") + 1] + matched_completions = [(cls, c) for (cls, c) in class_completions if c == cword_prefix] + elif len(cword_prefix) > 0: + matched_completions = [(cls, c) for (cls, c) in class_completions if c.startswith(cword_prefix)] + return matched_completions + + def inject_class_to_parser(self, cls): + """Add dummy arguments to our ArgumentParser for the traits of this class + + The argparse-based loader currently does not actually add any class traits to + the constructed ArgumentParser, only the flags & aliaes. In order to work nicely + with argcomplete's completers functionality, this method adds dummy arguments + of the form --Class.trait to the ArgumentParser instance. + + This method should be called selectively to reduce runtime overhead and to avoid + spamming options across all of Application.classes. + """ + try: + for traitname, trait in cls.class_traits(config=True).items(): + completer = trait.metadata.get("argcompleter") or getattr(trait, "argcompleter", None) + self._parser.add_argument( + f"--{cls.__name__}.{traitname}", + type=str, + help=trait.help, + # metavar=traitname, + ).completer = completer + argcomplete.debug(f"added --{cls.__name__}.{traitname}") + except AttributeError: + pass + + def _get_completions(self, comp_words, cword_prefix, *args): + """Overriden to dynamically append --Class.trait arguments if appropriate + + Warning: + This does not (currently) support completions of the form + --Class1.Class2.<...>.trait, although this is valid for traitlets. + Part of the reason is that we don't currently have a way to identify + which classes may be used with Class1 as a parent. + """ + prefix_chars = self._parser.prefix_chars + is_option = len(cword_prefix) > 0 and cword_prefix[0] in prefix_chars + if is_option: + # If we are currently completing an option, check if it could + # match with any of the --Class. completions. If there's exactly + # one matched class, then expand out the --Class.trait options. + matched_completions = self.match_class_completions(cword_prefix) + if len(matched_completions) == 1: + matched_cls = matched_completions[0][0] + self.inject_class_to_parser(matched_cls) + elif len(comp_words) > 0 and "." in comp_words[-1] and not is_option: + # If not an option, perform a hacky check to see if we are completing + # an argument for a --Class.trait option. + # TODO: the branch condition is wrong here for multiplicity="+", need to fix + matched_completions = self.match_class_completions(comp_words[-1]) + if matched_completions: + matched_cls = matched_completions[0][0] + self.inject_class_to_parser(matched_cls) + + return super()._get_completions(comp_words, cword_prefix, *args) + + def _get_option_completions(self, parser, cword_prefix): + """Overriden to add --Class. completions when appropriate""" + completions = super()._get_option_completions(parser, cword_prefix) + if cword_prefix.endswith("."): + return completions + + matched_completions = self.match_class_completions(cword_prefix) + if len(matched_completions) > 1: + completions.extend(opt for cls, opt in matched_completions) + # If there is exactly one match, we would expect it to have aleady + # been handled by the options dynamically added in _get_completions(). + # However, there could be edge cases, for example if the matched class + # has no configurable traits. + return completions + diff --git a/traitlets/config/loader.py b/traitlets/config/loader.py index 76d159a7..b6f3f52c 100644 --- a/traitlets/config/loader.py +++ b/traitlets/config/loader.py @@ -996,7 +996,7 @@ def _add_arguments(self, aliases, flags, classes): argparse_kwds["nargs"] = multiplicity argparse_traits[argname] = (trait, argparse_kwds) - for keys, (value, _) in flags.items(): + for keys, (value, fhelp) in flags.items(): if not isinstance(keys, tuple): keys = (keys,) for key in keys: @@ -1004,7 +1004,7 @@ def _add_arguments(self, aliases, flags, classes): alias_flags[aliases[key]] = value continue keys = ("-" + key, "--" + key) if len(key) == 1 else ("--" + key,) - paa(*keys, action=_FlagAction, flag=value) + paa(*keys, action=_FlagAction, flag=value, help=fhelp) for keys, traitname in aliases.items(): if not isinstance(keys, tuple): @@ -1045,7 +1045,7 @@ 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) - self.argcomplete() + self.argcomplete(classes) def _convert_to_config(self): """self.parsed_data->self.config, parse unrecognized extra args via KVLoader.""" @@ -1095,10 +1095,13 @@ def _handle_unrecognized_alias(self, arg: str): """ self.log.warning("Unrecognized alias: '%s', it will have no effect.", arg) - def argcomplete(self): + def argcomplete(self, classes: t.List[t.Any]): try: import argcomplete - argcomplete.autocomplete(self.parser) + from . import argcomplete_config + finder = argcomplete_config.ExtendedCompletionFinder() + finder.config_classes = classes # type: ignore + finder(self.parser) except ImportError: pass