From 523ec411286cb0bf4e3cd4a575decd66e42c6147 Mon Sep 17 00:00:00 2001 From: azhu Date: Fri, 9 Dec 2022 01:30:10 +1300 Subject: [PATCH] Custom argcomplete.CompletionFinder for class traits This custom 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. (This is currently done in a bit of a hacky manner.) 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. It also does not support nested class options like --Class1.Class2.trait. Example: ~/dev/traitlets$ examples/myapp.py --mode on --[TAB] --Application. --Foo. --debug --enable --help --j --mode --running --Bar. --MyApp. --disable --enabled --i --log_level --name ~/dev/traitlets$ examples/myapp.py --mode on --F[TAB] ~/dev/traitlets$ examples/myapp.py --mode on --Foo.[TAB] --Foo.i --Foo.j --Foo.mode --Foo.name ~/dev/traitlets$ examples/myapp.py --mode on --Foo.m[TAB] ~/dev/traitlets$ examples/myapp.py --mode on --Foo.mode [TAB] ~/dev/traitlets$ examples/myapp.py --mode on --Foo.mode o[TAB] off on other --- traitlets/config/argcomplete_config.py | 105 +++++++++++++++++++++++++ traitlets/config/loader.py | 13 +-- 2 files changed, 113 insertions(+), 5 deletions(-) create mode 100644 traitlets/config/argcomplete_config.py 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