221 lines
10 KiB
Python
221 lines
10 KiB
Python
|
"""Helper utilities for integrating argcomplete with traitlets"""
|
||
|
|
|
||
|
|
# Copyright (c) IPython Development Team.
|
||
|
|
# Distributed under the terms of the Modified BSD License.
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import typing as t
|
||
|
|
|
||
|
|
try:
|
||
|
|
import argcomplete
|
||
|
|
from argcomplete import CompletionFinder # type:ignore[attr-defined]
|
||
|
|
except ImportError:
|
||
|
|
# This module and its utility methods are written to not crash even
|
||
|
|
# if argcomplete is not installed.
|
||
|
|
class StubModule:
|
||
|
|
def __getattr__(self, attr: str) -> t.Any:
|
||
|
|
if not attr.startswith("__"):
|
||
|
|
raise ModuleNotFoundError("No module named 'argcomplete'")
|
||
|
|
raise AttributeError(f"argcomplete stub module has no attribute '{attr}'")
|
||
|
|
|
||
|
|
argcomplete = StubModule() # type:ignore[assignment]
|
||
|
|
CompletionFinder = object # type:ignore[assignment, misc]
|
||
|
|
|
||
|
|
|
||
|
|
def get_argcomplete_cwords() -> t.Optional[t.List[str]]:
|
||
|
|
"""Get current words prior to completion point
|
||
|
|
|
||
|
|
This is normally done in the `argcomplete.CompletionFinder` constructor,
|
||
|
|
but is exposed here to allow `traitlets` to follow dynamic code-paths such
|
||
|
|
as determining whether to evaluate a subcommand.
|
||
|
|
"""
|
||
|
|
if "_ARGCOMPLETE" not in os.environ:
|
||
|
|
return None
|
||
|
|
|
||
|
|
comp_line = os.environ["COMP_LINE"]
|
||
|
|
comp_point = int(os.environ["COMP_POINT"])
|
||
|
|
# argcomplete.debug("splitting COMP_LINE for:", comp_line, comp_point)
|
||
|
|
comp_words: t.List[str]
|
||
|
|
try:
|
||
|
|
(
|
||
|
|
cword_prequote,
|
||
|
|
cword_prefix,
|
||
|
|
cword_suffix,
|
||
|
|
comp_words,
|
||
|
|
last_wordbreak_pos,
|
||
|
|
) = argcomplete.split_line(comp_line, comp_point) # type:ignore[attr-defined,no-untyped-call]
|
||
|
|
except ModuleNotFoundError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# _ARGCOMPLETE is set by the shell script to tell us where comp_words
|
||
|
|
# should start, based on what we're completing.
|
||
|
|
# 1: <script> [args]
|
||
|
|
# 2: python <script> [args]
|
||
|
|
# 3: python -m <module> [args]
|
||
|
|
start = int(os.environ["_ARGCOMPLETE"]) - 1
|
||
|
|
comp_words = comp_words[start:]
|
||
|
|
|
||
|
|
# argcomplete.debug("prequote=", cword_prequote, "prefix=", cword_prefix, "suffix=", cword_suffix, "words=", comp_words, "last=", last_wordbreak_pos)
|
||
|
|
return comp_words # noqa: RET504
|
||
|
|
|
||
|
|
|
||
|
|
def increment_argcomplete_index() -> None:
|
||
|
|
"""Assumes ``$_ARGCOMPLETE`` is set and `argcomplete` is importable
|
||
|
|
|
||
|
|
Increment the index pointed to by ``$_ARGCOMPLETE``, which is used to
|
||
|
|
determine which word `argcomplete` should start evaluating the command-line.
|
||
|
|
This may be useful to "inform" `argcomplete` that we have already evaluated
|
||
|
|
the first word as a subcommand.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
os.environ["_ARGCOMPLETE"] = str(int(os.environ["_ARGCOMPLETE"]) + 1)
|
||
|
|
except Exception:
|
||
|
|
try:
|
||
|
|
argcomplete.debug("Unable to increment $_ARGCOMPLETE", os.environ["_ARGCOMPLETE"]) # type:ignore[attr-defined,no-untyped-call]
|
||
|
|
except (KeyError, ModuleNotFoundError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class ExtendedCompletionFinder(CompletionFinder):
|
||
|
|
"""An extension of CompletionFinder which dynamically completes class-trait based options
|
||
|
|
|
||
|
|
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,
|
||
|
|
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`,
|
||
|
|
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`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
_parser: argparse.ArgumentParser
|
||
|
|
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
|
||
|
|
|
||
|
|
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: t.Any) -> None:
|
||
|
|
"""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
|
||
|
|
)
|
||
|
|
multiplicity = trait.metadata.get("multiplicity")
|
||
|
|
self._parser.add_argument( # type: ignore[attr-defined]
|
||
|
|
f"--{cls.__name__}.{traitname}",
|
||
|
|
type=str,
|
||
|
|
help=trait.help,
|
||
|
|
nargs=multiplicity,
|
||
|
|
# metavar=traitname,
|
||
|
|
).completer = completer
|
||
|
|
# argcomplete.debug(f"added --{cls.__name__}.{traitname}")
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def _get_completions(
|
||
|
|
self, comp_words: t.List[str], cword_prefix: str, *args: t.Any
|
||
|
|
) -> t.List[str]:
|
||
|
|
"""Overridden 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.
|
||
|
|
|
||
|
|
Warning:
|
||
|
|
This is an internal method in CompletionFinder and so the API might
|
||
|
|
be subject to drift.
|
||
|
|
"""
|
||
|
|
# Try to identify if we are completing something related to --Class. for
|
||
|
|
# a known Class, if we are then add the Class config traits to our ArgumentParser.
|
||
|
|
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 an already present --Class.trait option. Search backwards
|
||
|
|
# for last option (based on last word starting with prefix_chars), and see
|
||
|
|
# if it is of the form --Class.trait. Note that if multiplicity="+", these
|
||
|
|
# arguments might conflict with positional arguments.
|
||
|
|
for prev_word in comp_words[::-1]:
|
||
|
|
if len(prev_word) > 0 and prev_word[0] in prefix_chars:
|
||
|
|
matched_completions = self.match_class_completions(prev_word)
|
||
|
|
if matched_completions:
|
||
|
|
matched_cls = matched_completions[0][0]
|
||
|
|
self.inject_class_to_parser(matched_cls)
|
||
|
|
break
|
||
|
|
|
||
|
|
completions: t.List[str]
|
||
|
|
completions = super()._get_completions(comp_words, cword_prefix, *args) # type:ignore[no-untyped-call]
|
||
|
|
|
||
|
|
# 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) # type:ignore[attr-defined,no-untyped-call]
|
||
|
|
completions.extend(subc for subc in self.subcommands if subc.startswith(cword_prefix))
|
||
|
|
|
||
|
|
return completions
|
||
|
|
|
||
|
|
def _get_option_completions(
|
||
|
|
self, parser: argparse.ArgumentParser, cword_prefix: str
|
||
|
|
) -> t.List[str]:
|
||
|
|
"""Overridden to add --Class. completions when appropriate"""
|
||
|
|
completions: t.List[str]
|
||
|
|
completions = super()._get_option_completions(parser, cword_prefix) # type:ignore[no-untyped-call]
|
||
|
|
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 already
|
||
|
|
# been handled by the options dynamically added in _get_completions().
|
||
|
|
# However, maybe there's an edge cases missed here, for example if the
|
||
|
|
# matched class has no configurable traits.
|
||
|
|
return completions
|