853 lines
31 KiB
Python
853 lines
31 KiB
Python
|
import logging
|
||
|
|
import re
|
||
|
|
from argparse import (
|
||
|
|
ONE_OR_MORE,
|
||
|
|
REMAINDER,
|
||
|
|
SUPPRESS,
|
||
|
|
ZERO_OR_MORE,
|
||
|
|
Action,
|
||
|
|
ArgumentParser,
|
||
|
|
_AppendAction,
|
||
|
|
_AppendConstAction,
|
||
|
|
_CountAction,
|
||
|
|
_HelpAction,
|
||
|
|
_StoreConstAction,
|
||
|
|
_VersionAction,
|
||
|
|
)
|
||
|
|
from collections import defaultdict
|
||
|
|
from functools import total_ordering
|
||
|
|
from itertools import starmap
|
||
|
|
from string import Template
|
||
|
|
from typing import Any, Dict, List
|
||
|
|
from typing import Optional as Opt
|
||
|
|
from typing import Union
|
||
|
|
|
||
|
|
# version detector. Precedence: installed dist, git, 'UNKNOWN'
|
||
|
|
try:
|
||
|
|
from ._dist_ver import __version__
|
||
|
|
except ImportError:
|
||
|
|
try:
|
||
|
|
from setuptools_scm import get_version
|
||
|
|
|
||
|
|
__version__ = get_version(root="..", relative_to=__file__)
|
||
|
|
except (ImportError, LookupError):
|
||
|
|
__version__ = "UNKNOWN"
|
||
|
|
__all__ = ["complete", "add_argument_to", "SUPPORTED_SHELLS", "FILE", "DIRECTORY", "DIR"]
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
SUPPORTED_SHELLS: List[str] = []
|
||
|
|
_SUPPORTED_COMPLETERS = {}
|
||
|
|
CHOICE_FUNCTIONS: Dict[str, Dict[str, str]] = {
|
||
|
|
"file": {"bash": "_shtab_compgen_files", "zsh": "_files", "tcsh": "f"},
|
||
|
|
"directory": {"bash": "_shtab_compgen_dirs", "zsh": "_files -/", "tcsh": "d"}}
|
||
|
|
FILE = CHOICE_FUNCTIONS["file"]
|
||
|
|
DIRECTORY = DIR = CHOICE_FUNCTIONS["directory"]
|
||
|
|
FLAG_OPTION = (
|
||
|
|
_StoreConstAction,
|
||
|
|
_HelpAction,
|
||
|
|
_VersionAction,
|
||
|
|
_AppendConstAction,
|
||
|
|
_CountAction,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class _ShtabPrintCompletionAction(Action):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
OPTION_END = _HelpAction, _VersionAction, _ShtabPrintCompletionAction
|
||
|
|
OPTION_MULTI = _AppendAction, _AppendConstAction, _CountAction
|
||
|
|
|
||
|
|
|
||
|
|
def mark_completer(shell):
|
||
|
|
def wrapper(func):
|
||
|
|
if shell not in SUPPORTED_SHELLS:
|
||
|
|
SUPPORTED_SHELLS.append(shell)
|
||
|
|
_SUPPORTED_COMPLETERS[shell] = func
|
||
|
|
return func
|
||
|
|
|
||
|
|
return wrapper
|
||
|
|
|
||
|
|
|
||
|
|
def get_completer(shell: str):
|
||
|
|
try:
|
||
|
|
return _SUPPORTED_COMPLETERS[shell]
|
||
|
|
except KeyError:
|
||
|
|
supported = ",".join(SUPPORTED_SHELLS)
|
||
|
|
raise NotImplementedError(f"shell ({shell}) must be in {supported}")
|
||
|
|
|
||
|
|
|
||
|
|
@total_ordering
|
||
|
|
class Choice:
|
||
|
|
"""
|
||
|
|
Placeholder to mark a special completion `<type>`.
|
||
|
|
|
||
|
|
>>> ArgumentParser.add_argument(..., choices=[Choice("<type>")])
|
||
|
|
"""
|
||
|
|
def __init__(self, choice_type: str, required: bool = False) -> None:
|
||
|
|
"""
|
||
|
|
See below for parameters.
|
||
|
|
|
||
|
|
choice_type : internal `type` name
|
||
|
|
required : controls result of comparison to empty strings
|
||
|
|
"""
|
||
|
|
self.required = required
|
||
|
|
self.type = choice_type
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
return self.type + ("" if self.required else "?")
|
||
|
|
|
||
|
|
def __cmp__(self, other: object) -> int:
|
||
|
|
if self.required:
|
||
|
|
return 0 if other else -1
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def __eq__(self, other: object) -> bool:
|
||
|
|
return self.__cmp__(other) == 0
|
||
|
|
|
||
|
|
def __lt__(self, other: object) -> bool:
|
||
|
|
return self.__cmp__(other) < 0
|
||
|
|
|
||
|
|
|
||
|
|
class Optional:
|
||
|
|
"""Example: `ArgumentParser.add_argument(..., choices=Optional.FILE)`."""
|
||
|
|
|
||
|
|
FILE = [Choice("file")]
|
||
|
|
DIR = DIRECTORY = [Choice("directory")]
|
||
|
|
|
||
|
|
|
||
|
|
class Required:
|
||
|
|
"""Example: `ArgumentParser.add_argument(..., choices=Required.FILE)`."""
|
||
|
|
|
||
|
|
FILE = [Choice("file", True)]
|
||
|
|
DIR = DIRECTORY = [Choice("directory", True)]
|
||
|
|
|
||
|
|
|
||
|
|
def complete2pattern(opt_complete, shell: str, choice_type2fn) -> str:
|
||
|
|
return (opt_complete.get(shell, "")
|
||
|
|
if isinstance(opt_complete, dict) else choice_type2fn[opt_complete])
|
||
|
|
|
||
|
|
|
||
|
|
def wordify(string: str) -> str:
|
||
|
|
"""Replace non-word chars [\\W] with underscores [_]"""
|
||
|
|
return re.sub("\\W", "_", string)
|
||
|
|
|
||
|
|
|
||
|
|
def get_public_subcommands(sub):
|
||
|
|
"""Get all the publicly-visible subcommands for a given subparser."""
|
||
|
|
public_parsers = {id(sub.choices[i.dest]) for i in sub._get_subactions()}
|
||
|
|
return {k for k, v in sub.choices.items() if id(v) in public_parsers}
|
||
|
|
|
||
|
|
|
||
|
|
def get_bash_commands(root_parser, root_prefix, choice_functions=None):
|
||
|
|
"""
|
||
|
|
Recursive subcommand parser traversal, returning lists of information on
|
||
|
|
commands (formatted for output to the completions script).
|
||
|
|
printing bash helper syntax.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
subparsers : list of subparsers for each parser
|
||
|
|
option_strings : list of options strings for each parser
|
||
|
|
compgens : list of shtab `.complete` functions corresponding to actions
|
||
|
|
choices : list of choices corresponding to actions
|
||
|
|
nargs : list of number of args allowed for each action (if not 0 or 1)
|
||
|
|
"""
|
||
|
|
choice_type2fn = {k: v["bash"] for k, v in CHOICE_FUNCTIONS.items()}
|
||
|
|
if choice_functions:
|
||
|
|
choice_type2fn.update(choice_functions)
|
||
|
|
|
||
|
|
def get_option_strings(parser):
|
||
|
|
"""Flattened list of all `parser`'s option strings."""
|
||
|
|
return sum(
|
||
|
|
(opt.option_strings for opt in parser._get_optional_actions() if opt.help != SUPPRESS),
|
||
|
|
[],
|
||
|
|
)
|
||
|
|
|
||
|
|
def recurse(parser, prefix):
|
||
|
|
"""recurse through subparsers, appending to the return lists"""
|
||
|
|
subparsers = []
|
||
|
|
option_strings = []
|
||
|
|
compgens = []
|
||
|
|
choices = []
|
||
|
|
nargs = []
|
||
|
|
|
||
|
|
# temp lists for recursion results
|
||
|
|
sub_subparsers = []
|
||
|
|
sub_option_strings = []
|
||
|
|
sub_compgens = []
|
||
|
|
sub_choices = []
|
||
|
|
sub_nargs = []
|
||
|
|
|
||
|
|
# positional arguments
|
||
|
|
discovered_subparsers = []
|
||
|
|
for i, positional in enumerate(parser._get_positional_actions()):
|
||
|
|
if positional.help == SUPPRESS:
|
||
|
|
continue
|
||
|
|
|
||
|
|
if hasattr(positional, "complete"):
|
||
|
|
# shtab `.complete = ...` functions
|
||
|
|
comp_pattern = complete2pattern(positional.complete, "bash", choice_type2fn)
|
||
|
|
compgens.append(f"{prefix}_pos_{i}_COMPGEN={comp_pattern}")
|
||
|
|
|
||
|
|
if positional.choices:
|
||
|
|
# choices (including subparsers & shtab `.complete` functions)
|
||
|
|
log.debug(f"choices:{prefix}:{sorted(positional.choices)}")
|
||
|
|
|
||
|
|
this_positional_choices = []
|
||
|
|
for choice in positional.choices:
|
||
|
|
if isinstance(choice, Choice):
|
||
|
|
# append special completion type to `compgens`
|
||
|
|
# NOTE: overrides `.complete` attribute
|
||
|
|
log.debug(f"Choice.{choice.type}:{prefix}:{positional.dest}")
|
||
|
|
compgens.append(f"{prefix}_pos_{i}_COMPGEN={choice_type2fn[choice.type]}")
|
||
|
|
elif isinstance(positional.choices, dict):
|
||
|
|
# subparser, so append to list of subparsers & recurse
|
||
|
|
log.debug("subcommand:%s", choice)
|
||
|
|
public_cmds = get_public_subcommands(positional)
|
||
|
|
if choice in public_cmds:
|
||
|
|
discovered_subparsers.append(str(choice))
|
||
|
|
this_positional_choices.append(str(choice))
|
||
|
|
(
|
||
|
|
new_subparsers,
|
||
|
|
new_option_strings,
|
||
|
|
new_compgens,
|
||
|
|
new_choices,
|
||
|
|
new_nargs,
|
||
|
|
) = recurse(
|
||
|
|
positional.choices[choice],
|
||
|
|
f"{prefix}_{wordify(choice)}",
|
||
|
|
)
|
||
|
|
sub_subparsers.extend(new_subparsers)
|
||
|
|
sub_option_strings.extend(new_option_strings)
|
||
|
|
sub_compgens.extend(new_compgens)
|
||
|
|
sub_choices.extend(new_choices)
|
||
|
|
sub_nargs.extend(new_nargs)
|
||
|
|
else:
|
||
|
|
log.debug("skip:subcommand:%s", choice)
|
||
|
|
else:
|
||
|
|
# simple choice
|
||
|
|
this_positional_choices.append(str(choice))
|
||
|
|
|
||
|
|
if this_positional_choices:
|
||
|
|
choices_str = "' '".join(this_positional_choices)
|
||
|
|
choices.append(f"{prefix}_pos_{i}_choices=('{choices_str}')")
|
||
|
|
|
||
|
|
# skip default `nargs` values
|
||
|
|
if positional.nargs not in (None, "1", "?"):
|
||
|
|
nargs.append(f"{prefix}_pos_{i}_nargs={positional.nargs}")
|
||
|
|
|
||
|
|
if discovered_subparsers:
|
||
|
|
subparsers_str = "' '".join(discovered_subparsers)
|
||
|
|
subparsers.append(f"{prefix}_subparsers=('{subparsers_str}')")
|
||
|
|
log.debug(f"subcommands:{prefix}:{discovered_subparsers}")
|
||
|
|
|
||
|
|
# optional arguments
|
||
|
|
options_strings_str = "' '".join(get_option_strings(parser))
|
||
|
|
option_strings.append(f"{prefix}_option_strings=('{options_strings_str}')")
|
||
|
|
for optional in parser._get_optional_actions():
|
||
|
|
if optional == SUPPRESS:
|
||
|
|
continue
|
||
|
|
|
||
|
|
for option_string in optional.option_strings:
|
||
|
|
if hasattr(optional, "complete"):
|
||
|
|
# shtab `.complete = ...` functions
|
||
|
|
comp_pattern_str = complete2pattern(optional.complete, "bash", choice_type2fn)
|
||
|
|
compgens.append(
|
||
|
|
f"{prefix}_{wordify(option_string)}_COMPGEN={comp_pattern_str}")
|
||
|
|
|
||
|
|
if optional.choices:
|
||
|
|
# choices (including shtab `.complete` functions)
|
||
|
|
this_optional_choices = []
|
||
|
|
for choice in optional.choices:
|
||
|
|
# append special completion type to `compgens`
|
||
|
|
# NOTE: overrides `.complete` attribute
|
||
|
|
if isinstance(choice, Choice):
|
||
|
|
log.debug(f"Choice.{choice.type}:{prefix}:{optional.dest}")
|
||
|
|
func_str = choice_type2fn[choice.type]
|
||
|
|
compgens.append(
|
||
|
|
f"{prefix}_{wordify(option_string)}_COMPGEN={func_str}")
|
||
|
|
else:
|
||
|
|
# simple choice
|
||
|
|
this_optional_choices.append(str(choice))
|
||
|
|
|
||
|
|
if this_optional_choices:
|
||
|
|
this_choices_str = "' '".join(this_optional_choices)
|
||
|
|
choices.append(
|
||
|
|
f"{prefix}_{wordify(option_string)}_choices=('{this_choices_str}')")
|
||
|
|
|
||
|
|
# Check for nargs.
|
||
|
|
if optional.nargs is not None and optional.nargs != 1:
|
||
|
|
nargs.append(f"{prefix}_{wordify(option_string)}_nargs={optional.nargs}")
|
||
|
|
|
||
|
|
# append recursion results
|
||
|
|
subparsers.extend(sub_subparsers)
|
||
|
|
option_strings.extend(sub_option_strings)
|
||
|
|
compgens.extend(sub_compgens)
|
||
|
|
choices.extend(sub_choices)
|
||
|
|
nargs.extend(sub_nargs)
|
||
|
|
|
||
|
|
return subparsers, option_strings, compgens, choices, nargs
|
||
|
|
|
||
|
|
return recurse(root_parser, root_prefix)
|
||
|
|
|
||
|
|
|
||
|
|
@mark_completer("bash")
|
||
|
|
def complete_bash(parser, root_prefix=None, preamble="", choice_functions=None):
|
||
|
|
"""
|
||
|
|
Returns bash syntax autocompletion script.
|
||
|
|
|
||
|
|
See `complete` for arguments.
|
||
|
|
"""
|
||
|
|
root_prefix = wordify(f"_shtab_{root_prefix or parser.prog}")
|
||
|
|
subparsers, option_strings, compgens, choices, nargs = get_bash_commands(
|
||
|
|
parser, root_prefix, choice_functions=choice_functions)
|
||
|
|
|
||
|
|
# References:
|
||
|
|
# - https://www.gnu.org/software/bash/manual/html_node/
|
||
|
|
# Programmable-Completion.html
|
||
|
|
# - https://opensource.com/article/18/3/creating-bash-completion-script
|
||
|
|
# - https://stackoverflow.com/questions/12933362
|
||
|
|
return Template("""\
|
||
|
|
# AUTOMATICALLY GENERATED by `shtab`
|
||
|
|
|
||
|
|
${subparsers}
|
||
|
|
|
||
|
|
${option_strings}
|
||
|
|
|
||
|
|
${compgens}
|
||
|
|
|
||
|
|
${choices}
|
||
|
|
|
||
|
|
${nargs}
|
||
|
|
|
||
|
|
${preamble}
|
||
|
|
# $1=COMP_WORDS[1]
|
||
|
|
_shtab_compgen_files() {
|
||
|
|
compgen -f -- $1 # files
|
||
|
|
}
|
||
|
|
|
||
|
|
# $1=COMP_WORDS[1]
|
||
|
|
_shtab_compgen_dirs() {
|
||
|
|
compgen -d -- $1 # recurse into subdirs
|
||
|
|
}
|
||
|
|
|
||
|
|
# $1=COMP_WORDS[1]
|
||
|
|
_shtab_replace_nonword() {
|
||
|
|
echo "${1//[^[:word:]]/_}"
|
||
|
|
}
|
||
|
|
|
||
|
|
# set default values (called for the initial parser & any subparsers)
|
||
|
|
_set_parser_defaults() {
|
||
|
|
local subparsers_var="${prefix}_subparsers[@]"
|
||
|
|
sub_parsers=${!subparsers_var-}
|
||
|
|
|
||
|
|
local current_option_strings_var="${prefix}_option_strings[@]"
|
||
|
|
current_option_strings=${!current_option_strings_var}
|
||
|
|
|
||
|
|
completed_positional_actions=0
|
||
|
|
|
||
|
|
_set_new_action "pos_${completed_positional_actions}" true
|
||
|
|
}
|
||
|
|
|
||
|
|
# $1=action identifier
|
||
|
|
# $2=positional action (bool)
|
||
|
|
# set all identifiers for an action's parameters
|
||
|
|
_set_new_action() {
|
||
|
|
current_action="${prefix}_$(_shtab_replace_nonword $1)"
|
||
|
|
|
||
|
|
local current_action_compgen_var=${current_action}_COMPGEN
|
||
|
|
current_action_compgen="${!current_action_compgen_var-}"
|
||
|
|
|
||
|
|
local current_action_choices_var="${current_action}_choices[@]"
|
||
|
|
current_action_choices="${!current_action_choices_var-}"
|
||
|
|
|
||
|
|
local current_action_nargs_var="${current_action}_nargs"
|
||
|
|
if [ -n "${!current_action_nargs_var-}" ]; then
|
||
|
|
current_action_nargs="${!current_action_nargs_var}"
|
||
|
|
else
|
||
|
|
current_action_nargs=1
|
||
|
|
fi
|
||
|
|
|
||
|
|
current_action_args_start_index=$(( $word_index + 1 - $pos_only ))
|
||
|
|
|
||
|
|
current_action_is_positional=$2
|
||
|
|
}
|
||
|
|
|
||
|
|
# Notes:
|
||
|
|
# `COMPREPLY`: what will be rendered after completion is triggered
|
||
|
|
# `completing_word`: currently typed word to generate completions for
|
||
|
|
# `${!var}`: evaluates the content of `var` and expand its content as a variable
|
||
|
|
# hello="world"
|
||
|
|
# x="hello"
|
||
|
|
# ${!x} -> ${hello} -> "world"
|
||
|
|
${root_prefix}() {
|
||
|
|
local completing_word="${COMP_WORDS[COMP_CWORD]}"
|
||
|
|
local previous_word="${COMP_WORDS[COMP_CWORD-1]}"
|
||
|
|
local completed_positional_actions
|
||
|
|
local current_action
|
||
|
|
local current_action_args_start_index
|
||
|
|
local current_action_choices
|
||
|
|
local current_action_compgen
|
||
|
|
local current_action_is_positional
|
||
|
|
local current_action_nargs
|
||
|
|
local current_option_strings
|
||
|
|
local sub_parsers
|
||
|
|
COMPREPLY=()
|
||
|
|
|
||
|
|
local prefix=${root_prefix}
|
||
|
|
local word_index=0
|
||
|
|
local pos_only=0 # "--" delimeter not encountered yet
|
||
|
|
_set_parser_defaults
|
||
|
|
word_index=1
|
||
|
|
|
||
|
|
# determine what arguments are appropriate for the current state
|
||
|
|
# of the arg parser
|
||
|
|
while [ $word_index -ne $COMP_CWORD ]; do
|
||
|
|
local this_word="${COMP_WORDS[$word_index]}"
|
||
|
|
|
||
|
|
if [[ $pos_only = 1 || " $this_word " != " -- " ]]; then
|
||
|
|
if [[ -n $sub_parsers && " ${sub_parsers[@]} " == *" ${this_word} "* ]]; then
|
||
|
|
# valid subcommand: add it to the prefix & reset the current action
|
||
|
|
prefix="${prefix}_$(_shtab_replace_nonword $this_word)"
|
||
|
|
_set_parser_defaults
|
||
|
|
fi
|
||
|
|
|
||
|
|
if [[ " ${current_option_strings[@]} " == *" ${this_word} "* ]]; then
|
||
|
|
# a new action should be acquired (due to recognised option string or
|
||
|
|
# no more input expected from current action);
|
||
|
|
# the next positional action can fill in here
|
||
|
|
_set_new_action $this_word false
|
||
|
|
fi
|
||
|
|
|
||
|
|
if [[ "$current_action_nargs" != "*" ]] && \\
|
||
|
|
[[ "$current_action_nargs" != "+" ]] && \\
|
||
|
|
[[ "$current_action_nargs" != *"..." ]] && \\
|
||
|
|
(( $word_index + 1 - $current_action_args_start_index - $pos_only >= \\
|
||
|
|
$current_action_nargs )); then
|
||
|
|
$current_action_is_positional && let "completed_positional_actions += 1"
|
||
|
|
_set_new_action "pos_${completed_positional_actions}" true
|
||
|
|
fi
|
||
|
|
else
|
||
|
|
pos_only=1 # "--" delimeter encountered
|
||
|
|
fi
|
||
|
|
|
||
|
|
let "word_index+=1"
|
||
|
|
done
|
||
|
|
|
||
|
|
# Generate the completions
|
||
|
|
|
||
|
|
if [[ $pos_only = 0 && "${completing_word}" == -* ]]; then
|
||
|
|
# optional argument started: use option strings
|
||
|
|
COMPREPLY=( $(compgen -W "${current_option_strings[*]}" -- "${completing_word}") )
|
||
|
|
elif [[ "${previous_word}" == ">" || "${previous_word}" == ">>" ||
|
||
|
|
"${previous_word}" =~ ^[12]">" || "${previous_word}" =~ ^[12]">>" ]]; then
|
||
|
|
# handle redirection operators
|
||
|
|
COMPREPLY=( $(compgen -f -- "${completing_word}") )
|
||
|
|
else
|
||
|
|
# use choices & compgen
|
||
|
|
local IFS=$'\\n' # items may contain spaces, so delimit using newline
|
||
|
|
COMPREPLY=( $([ -n "${current_action_compgen}" ] \\
|
||
|
|
&& "${current_action_compgen}" "${completing_word}") )
|
||
|
|
unset IFS
|
||
|
|
COMPREPLY+=( $(compgen -W "${current_action_choices[*]}" -- "${completing_word}") )
|
||
|
|
fi
|
||
|
|
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
|
||
|
|
complete -o filenames -F ${root_prefix} ${prog}""").safe_substitute(
|
||
|
|
subparsers="\n".join(subparsers),
|
||
|
|
option_strings="\n".join(option_strings),
|
||
|
|
compgens="\n".join(compgens),
|
||
|
|
choices="\n".join(choices),
|
||
|
|
nargs="\n".join(nargs),
|
||
|
|
preamble=("\n# Custom Preamble\n" + preamble +
|
||
|
|
"\n# End Custom Preamble\n" if preamble else ""),
|
||
|
|
root_prefix=root_prefix,
|
||
|
|
prog=parser.prog,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def escape_zsh(string):
|
||
|
|
# excessive but safe
|
||
|
|
return re.sub(r"([^\w\s.,()-])", r"\\\1", str(string))
|
||
|
|
|
||
|
|
|
||
|
|
@mark_completer("zsh")
|
||
|
|
def complete_zsh(parser, root_prefix=None, preamble="", choice_functions=None):
|
||
|
|
"""
|
||
|
|
Returns zsh syntax autocompletion script.
|
||
|
|
|
||
|
|
See `complete` for arguments.
|
||
|
|
"""
|
||
|
|
prog = parser.prog
|
||
|
|
root_prefix = wordify(f"_shtab_{root_prefix or prog}")
|
||
|
|
|
||
|
|
choice_type2fn = {k: v["zsh"] for k, v in CHOICE_FUNCTIONS.items()}
|
||
|
|
if choice_functions:
|
||
|
|
choice_type2fn.update(choice_functions)
|
||
|
|
|
||
|
|
def is_opt_end(opt):
|
||
|
|
return isinstance(opt, OPTION_END) or opt.nargs == REMAINDER
|
||
|
|
|
||
|
|
def is_opt_multiline(opt):
|
||
|
|
return isinstance(opt, OPTION_MULTI)
|
||
|
|
|
||
|
|
def format_optional(opt, parser):
|
||
|
|
get_help = parser._get_formatter()._expand_help
|
||
|
|
return (('{nargs}{options}"[{help}]"' if isinstance(
|
||
|
|
opt, FLAG_OPTION) else '{nargs}{options}"[{help}]:{dest}:{pattern}"').format(
|
||
|
|
nargs=('"(- : *)"' if is_opt_end(opt) else '"*"' if is_opt_multiline(opt) else ""),
|
||
|
|
options=("{{{}}}".format(",".join(opt.option_strings)) if len(opt.option_strings)
|
||
|
|
> 1 else '"{}"'.format("".join(opt.option_strings))),
|
||
|
|
help=escape_zsh(get_help(opt) if opt.help else ""),
|
||
|
|
dest=opt.dest,
|
||
|
|
pattern=complete2pattern(opt.complete, "zsh", choice_type2fn) if hasattr(
|
||
|
|
opt, "complete") else
|
||
|
|
(choice_type2fn[opt.choices[0].type] if isinstance(opt.choices[0], Choice) else
|
||
|
|
"({})".format(" ".join(map(str, opt.choices)))) if opt.choices else "",
|
||
|
|
).replace('""', ""))
|
||
|
|
|
||
|
|
def format_positional(opt, parser):
|
||
|
|
get_help = parser._get_formatter()._expand_help
|
||
|
|
return '"{nargs}:{help}:{pattern}"'.format(
|
||
|
|
nargs={ONE_OR_MORE: "(*)", ZERO_OR_MORE: "(*):", REMAINDER: "(-)*"}.get(opt.nargs, ""),
|
||
|
|
help=escape_zsh((get_help(opt) if opt.help else opt.dest).strip().split("\n")[0]),
|
||
|
|
pattern=complete2pattern(opt.complete, "zsh", choice_type2fn) if hasattr(
|
||
|
|
opt, "complete") else
|
||
|
|
(choice_type2fn[opt.choices[0].type] if isinstance(opt.choices[0], Choice) else
|
||
|
|
"({})".format(" ".join(map(str, opt.choices)))) if opt.choices else "",
|
||
|
|
)
|
||
|
|
|
||
|
|
# {cmd: {"help": help, "arguments": [arguments]}}
|
||
|
|
all_commands = {
|
||
|
|
root_prefix: {
|
||
|
|
"cmd": prog, "arguments": [
|
||
|
|
format_optional(opt, parser)
|
||
|
|
for opt in parser._get_optional_actions() if opt.help != SUPPRESS] + [
|
||
|
|
format_positional(opt, parser) for opt in parser._get_positional_actions()
|
||
|
|
if opt.help != SUPPRESS and opt.choices is None],
|
||
|
|
"help": (parser.description
|
||
|
|
or "").strip().split("\n")[0], "commands": [], "paths": []}}
|
||
|
|
|
||
|
|
def recurse(parser, prefix, paths=None):
|
||
|
|
paths = paths or []
|
||
|
|
subcmds = []
|
||
|
|
for sub in parser._get_positional_actions():
|
||
|
|
if sub.help == SUPPRESS or not sub.choices:
|
||
|
|
continue
|
||
|
|
if not sub.choices or not isinstance(sub.choices, dict):
|
||
|
|
# positional argument
|
||
|
|
all_commands[prefix]["arguments"].append(format_positional(sub, parser))
|
||
|
|
else: # subparser
|
||
|
|
log.debug(f"choices:{prefix}:{sorted(sub.choices)}")
|
||
|
|
public_cmds = get_public_subcommands(sub)
|
||
|
|
for cmd, subparser in sub.choices.items():
|
||
|
|
if cmd not in public_cmds:
|
||
|
|
log.debug("skip:subcommand:%s", cmd)
|
||
|
|
continue
|
||
|
|
log.debug("subcommand:%s", cmd)
|
||
|
|
|
||
|
|
# optionals
|
||
|
|
arguments = [
|
||
|
|
format_optional(opt, parser) for opt in subparser._get_optional_actions()
|
||
|
|
if opt.help != SUPPRESS]
|
||
|
|
|
||
|
|
# positionals
|
||
|
|
arguments.extend(
|
||
|
|
format_positional(opt, parser)
|
||
|
|
for opt in subparser._get_positional_actions()
|
||
|
|
if not isinstance(opt.choices, dict) if opt.help != SUPPRESS)
|
||
|
|
|
||
|
|
# help text
|
||
|
|
formatter = subparser._get_formatter()
|
||
|
|
backup_width = formatter._width
|
||
|
|
formatter._width = 1234567 # large number to effectively disable wrapping
|
||
|
|
desc = formatter._format_text(subparser.description or "").strip()
|
||
|
|
formatter._width = backup_width
|
||
|
|
|
||
|
|
new_pref = f"{prefix}_{wordify(cmd)}"
|
||
|
|
options = all_commands[new_pref] = {
|
||
|
|
"cmd": cmd, "help": desc.split("\n")[0], "arguments": arguments,
|
||
|
|
"paths": [*paths, cmd]}
|
||
|
|
new_subcmds = recurse(subparser, new_pref, [*paths, cmd])
|
||
|
|
options["commands"] = {
|
||
|
|
all_commands[pref]["cmd"]: all_commands[pref]
|
||
|
|
for pref in new_subcmds if pref in all_commands}
|
||
|
|
subcmds.extend([*new_subcmds, new_pref])
|
||
|
|
log.debug("subcommands:%s:%s", cmd, options)
|
||
|
|
return subcmds
|
||
|
|
|
||
|
|
recurse(parser, root_prefix)
|
||
|
|
all_commands[root_prefix]["commands"] = {
|
||
|
|
options["cmd"]: options
|
||
|
|
for prefix, options in sorted(all_commands.items())
|
||
|
|
if len(options.get("paths", [])) < 2 and prefix != root_prefix}
|
||
|
|
subcommands = {
|
||
|
|
prefix: options
|
||
|
|
for prefix, options in all_commands.items() if options.get("commands")}
|
||
|
|
subcommands.setdefault(root_prefix, all_commands[root_prefix])
|
||
|
|
log.debug("subcommands:%s:%s", root_prefix, sorted(all_commands))
|
||
|
|
|
||
|
|
def command_case(prefix, options):
|
||
|
|
name = options["cmd"]
|
||
|
|
commands = options["commands"]
|
||
|
|
case_fmt_on_no_sub = """{name}) _arguments -C -s ${prefix}_{name_wordify}_options ;;"""
|
||
|
|
case_fmt_on_sub = """{name}) {prefix}_{name_wordify} ;;"""
|
||
|
|
|
||
|
|
cases = []
|
||
|
|
for _, options in sorted(commands.items()):
|
||
|
|
fmt = case_fmt_on_sub if options.get("commands") else case_fmt_on_no_sub
|
||
|
|
cases.append(
|
||
|
|
fmt.format(name=options["cmd"], name_wordify=wordify(options["cmd"]),
|
||
|
|
prefix=prefix))
|
||
|
|
cases = "\n\t".expandtabs(8).join(cases)
|
||
|
|
|
||
|
|
return f"""\
|
||
|
|
{prefix}() {{
|
||
|
|
local context state line curcontext="$curcontext" one_or_more='(-)*' remainder='(*)'
|
||
|
|
|
||
|
|
if ((${{{prefix}_options[(I)${{(q)one_or_more}}*]}} + ${{{prefix}_options[(I)${{(q)remainder}}*]}} == 0)); then # noqa: E501
|
||
|
|
{prefix}_options+=(': :{prefix}_commands' '*::: :->{name}')
|
||
|
|
fi
|
||
|
|
_arguments -C -s ${prefix}_options
|
||
|
|
|
||
|
|
case $state in
|
||
|
|
{name})
|
||
|
|
words=($line[1] "${{words[@]}}")
|
||
|
|
(( CURRENT += 1 ))
|
||
|
|
curcontext="${{curcontext%:*:*}}:{prefix}-$line[1]:"
|
||
|
|
case $line[1] in
|
||
|
|
{cases}
|
||
|
|
esac
|
||
|
|
esac
|
||
|
|
}}
|
||
|
|
"""
|
||
|
|
|
||
|
|
def command_option(prefix, options):
|
||
|
|
arguments = "\n ".join(options["arguments"])
|
||
|
|
return f"""\
|
||
|
|
{prefix}_options=(
|
||
|
|
{arguments}
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
|
||
|
|
def command_list(prefix, options):
|
||
|
|
name = " ".join([prog, *options["paths"]])
|
||
|
|
commands = "\n ".join(f'"{escape_zsh(cmd)}:{escape_zsh(opt["help"])}"'
|
||
|
|
for cmd, opt in sorted(options["commands"].items()))
|
||
|
|
return f"""
|
||
|
|
{prefix}_commands() {{
|
||
|
|
local _commands=(
|
||
|
|
{commands}
|
||
|
|
)
|
||
|
|
_describe '{name} commands' _commands
|
||
|
|
}}"""
|
||
|
|
|
||
|
|
preamble = (f"""\
|
||
|
|
# Custom Preamble
|
||
|
|
{preamble.rstrip()}
|
||
|
|
|
||
|
|
# End Custom Preamble
|
||
|
|
""" if preamble else "")
|
||
|
|
# References:
|
||
|
|
# - https://github.com/zsh-users/zsh-completions
|
||
|
|
# - http://zsh.sourceforge.net/Doc/Release/Completion-System.html
|
||
|
|
# - https://mads-hartmann.com/2017/08/06/
|
||
|
|
# writing-zsh-completion-scripts.html
|
||
|
|
# - http://www.linux-mag.com/id/1106/
|
||
|
|
return Template("""\
|
||
|
|
#compdef ${prog}
|
||
|
|
|
||
|
|
# AUTOMATICALLY GENERATED by `shtab`
|
||
|
|
|
||
|
|
${command_commands}
|
||
|
|
|
||
|
|
${command_options}
|
||
|
|
|
||
|
|
${command_cases}
|
||
|
|
${preamble}
|
||
|
|
|
||
|
|
typeset -A opt_args
|
||
|
|
|
||
|
|
if [[ $zsh_eval_context[-1] == eval ]]; then
|
||
|
|
# eval/source/. command, register function for later
|
||
|
|
compdef ${root_prefix} -N ${prog}
|
||
|
|
else
|
||
|
|
# autoload from fpath, call function directly
|
||
|
|
${root_prefix} "$@\"
|
||
|
|
fi
|
||
|
|
""").safe_substitute(
|
||
|
|
prog=prog,
|
||
|
|
root_prefix=root_prefix,
|
||
|
|
command_cases="\n".join(starmap(command_case, sorted(subcommands.items()))),
|
||
|
|
command_commands="\n".join(starmap(command_list, sorted(subcommands.items()))),
|
||
|
|
command_options="\n".join(starmap(command_option, sorted(all_commands.items()))),
|
||
|
|
preamble=preamble,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@mark_completer("tcsh")
|
||
|
|
def complete_tcsh(parser, root_prefix=None, preamble="", choice_functions=None):
|
||
|
|
"""
|
||
|
|
Return tcsh syntax autocompletion script.
|
||
|
|
|
||
|
|
root_prefix:
|
||
|
|
ignored (tcsh has no support for functions)
|
||
|
|
|
||
|
|
See `complete` for other arguments.
|
||
|
|
"""
|
||
|
|
optionals_single = set()
|
||
|
|
optionals_double = set()
|
||
|
|
specials = []
|
||
|
|
index_choices = defaultdict(dict)
|
||
|
|
|
||
|
|
choice_type2fn = {k: v["tcsh"] for k, v in CHOICE_FUNCTIONS.items()}
|
||
|
|
if choice_functions:
|
||
|
|
choice_type2fn.update(choice_functions)
|
||
|
|
|
||
|
|
def get_specials(arg, arg_type, arg_sel):
|
||
|
|
if arg.choices:
|
||
|
|
choice_strs = ' '.join(map(str, arg.choices))
|
||
|
|
yield f"'{arg_type}/{arg_sel}/({choice_strs})/'"
|
||
|
|
elif hasattr(arg, 'complete'):
|
||
|
|
complete_fn = complete2pattern(arg.complete, 'tcsh', choice_type2fn)
|
||
|
|
if complete_fn:
|
||
|
|
yield f"'{arg_type}/{arg_sel}/{complete_fn}/'"
|
||
|
|
|
||
|
|
def recurse_parser(cparser, positional_idx, requirements=None):
|
||
|
|
log_prefix = "| " * positional_idx
|
||
|
|
log.debug("%sParser @ %d", log_prefix, positional_idx)
|
||
|
|
if requirements:
|
||
|
|
log.debug("%s- Requires: %s", log_prefix, " ".join(requirements))
|
||
|
|
else:
|
||
|
|
requirements = []
|
||
|
|
|
||
|
|
for optional in cparser._get_optional_actions():
|
||
|
|
log.debug("%s| Optional: %s", log_prefix, optional.dest)
|
||
|
|
if optional.help != SUPPRESS:
|
||
|
|
# Mingle all optional arguments for all subparsers
|
||
|
|
for optional_str in optional.option_strings:
|
||
|
|
log.debug("%s| | %s", log_prefix, optional_str)
|
||
|
|
if optional_str.startswith('--'):
|
||
|
|
optionals_double.add(optional_str[2:])
|
||
|
|
elif optional_str.startswith('-'):
|
||
|
|
optionals_single.add(optional_str[1:])
|
||
|
|
specials.extend(get_specials(optional, 'n', optional_str))
|
||
|
|
|
||
|
|
for positional in cparser._get_positional_actions():
|
||
|
|
if positional.help != SUPPRESS:
|
||
|
|
positional_idx += 1
|
||
|
|
log.debug("%s| Positional #%d: %s", log_prefix, positional_idx, positional.dest)
|
||
|
|
index_choices[positional_idx][tuple(requirements)] = positional
|
||
|
|
if not requirements and isinstance(positional.choices, dict):
|
||
|
|
for subcmd, subparser in positional.choices.items():
|
||
|
|
log.debug("%s| | SubParser: %s", log_prefix, subcmd)
|
||
|
|
recurse_parser(subparser, positional_idx, requirements + [subcmd])
|
||
|
|
|
||
|
|
recurse_parser(parser, 0)
|
||
|
|
|
||
|
|
for idx, ndict in index_choices.items():
|
||
|
|
if len(ndict) == 1:
|
||
|
|
# Single choice, no requirements
|
||
|
|
arg = list(ndict.values())[0]
|
||
|
|
specials.extend(get_specials(arg, 'p', str(idx)))
|
||
|
|
else:
|
||
|
|
# Multiple requirements
|
||
|
|
nlist = []
|
||
|
|
for nn, arg in ndict.items():
|
||
|
|
if arg.choices:
|
||
|
|
checks = [f'[ "$cmd[{iidx}]" == "{n}" ]' for iidx, n in enumerate(nn, start=2)]
|
||
|
|
choices_str = "' '".join(arg.choices)
|
||
|
|
checks_str = ' && '.join(checks + [f"echo '{choices_str}'"])
|
||
|
|
nlist.append(f"( {checks_str} || false )")
|
||
|
|
# Ugly hack
|
||
|
|
nlist_str = ' || '.join(nlist)
|
||
|
|
specials.append(f"'p@{str(idx)}@`set cmd=($COMMAND_LINE); {nlist_str}`@'")
|
||
|
|
|
||
|
|
if optionals_double:
|
||
|
|
if optionals_single:
|
||
|
|
optionals_single.add('-')
|
||
|
|
else:
|
||
|
|
# Don't add a space after completing "--" from "-"
|
||
|
|
optionals_single = ('-', '-')
|
||
|
|
|
||
|
|
return Template("""\
|
||
|
|
# AUTOMATICALLY GENERATED by `shtab`
|
||
|
|
|
||
|
|
${preamble}
|
||
|
|
|
||
|
|
complete ${prog} \\
|
||
|
|
'c/--/(${optionals_double_str})/' \\
|
||
|
|
'c/-/(${optionals_single_str})/' \\
|
||
|
|
${optionals_special_str} \\
|
||
|
|
'p/*/()/'""").safe_substitute(
|
||
|
|
preamble=("\n# Custom Preamble\n" + preamble +
|
||
|
|
"\n# End Custom Preamble\n" if preamble else ""), root_prefix=root_prefix,
|
||
|
|
prog=parser.prog, optionals_double_str=' '.join(sorted(optionals_double)),
|
||
|
|
optionals_single_str=' '.join(sorted(optionals_single)),
|
||
|
|
optionals_special_str=' \\\n '.join(specials))
|
||
|
|
|
||
|
|
|
||
|
|
def complete(parser: ArgumentParser, shell: str = "bash", root_prefix: Opt[str] = None,
|
||
|
|
preamble: Union[str, Dict[str, str]] = "", choice_functions: Opt[Any] = None) -> str:
|
||
|
|
"""
|
||
|
|
shell:
|
||
|
|
bash/zsh/tcsh
|
||
|
|
root_prefix:
|
||
|
|
prefix for shell functions to avoid clashes (default: "_{parser.prog}")
|
||
|
|
preamble:
|
||
|
|
mapping shell to text to prepend to generated script
|
||
|
|
(e.g. `{"bash": "_myprog_custom_function(){ echo hello }"}`)
|
||
|
|
choice_functions:
|
||
|
|
*deprecated*
|
||
|
|
|
||
|
|
N.B. `parser.add_argument().complete = ...` can be used to define custom
|
||
|
|
completions (e.g. filenames). See <../examples/pathcomplete.py>.
|
||
|
|
"""
|
||
|
|
if isinstance(preamble, dict):
|
||
|
|
preamble = preamble.get(shell, "")
|
||
|
|
completer = get_completer(shell)
|
||
|
|
return completer(
|
||
|
|
parser,
|
||
|
|
root_prefix=root_prefix,
|
||
|
|
preamble=preamble,
|
||
|
|
choice_functions=choice_functions,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def completion_action(parent: Opt[ArgumentParser] = None, preamble: Union[str, Dict[str,
|
||
|
|
str]] = ""):
|
||
|
|
class PrintCompletionAction(_ShtabPrintCompletionAction):
|
||
|
|
def __call__(self, parser, namespace, values, option_string=None):
|
||
|
|
print(complete(parent or parser, values, preamble=preamble))
|
||
|
|
parser.exit(0)
|
||
|
|
|
||
|
|
return PrintCompletionAction
|
||
|
|
|
||
|
|
|
||
|
|
def add_argument_to(
|
||
|
|
parser: ArgumentParser,
|
||
|
|
option_string: Union[str, List[str]] = "--print-completion",
|
||
|
|
help: str = "print shell completion script",
|
||
|
|
parent: Opt[ArgumentParser] = None,
|
||
|
|
preamble: Union[str, Dict[str, str]] = "",
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
option_string:
|
||
|
|
iff positional (no `-` prefix) then `parser` is assumed to actually be
|
||
|
|
a subparser (subcommand mode)
|
||
|
|
parent:
|
||
|
|
required in subcommand mode
|
||
|
|
"""
|
||
|
|
if isinstance(option_string, str):
|
||
|
|
option_string = [option_string]
|
||
|
|
kwargs = {
|
||
|
|
"choices": SUPPORTED_SHELLS, "default": None, "help": help,
|
||
|
|
"action": completion_action(parent, preamble)}
|
||
|
|
if option_string[0][0] != "-": # subparser mode
|
||
|
|
kwargs.update(default=SUPPORTED_SHELLS[0], nargs="?")
|
||
|
|
assert parent is not None, "subcommand mode: parent required"
|
||
|
|
parser.add_argument(*option_string, **kwargs)
|
||
|
|
return parser
|