_bashcomplete.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  1. import copy
  2. import os
  3. import re
  4. from .utils import echo
  5. from .parser import split_arg_string
  6. from .core import MultiCommand, Option, Argument
  7. from .types import Choice
  8. try:
  9. from collections import abc
  10. except ImportError:
  11. import collections as abc
  12. WORDBREAK = '='
  13. # Note, only BASH version 4.4 and later have the nosort option.
  14. COMPLETION_SCRIPT_BASH = '''
  15. %(complete_func)s() {
  16. local IFS=$'\n'
  17. COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \\
  18. COMP_CWORD=$COMP_CWORD \\
  19. %(autocomplete_var)s=complete $1 ) )
  20. return 0
  21. }
  22. %(complete_func)setup() {
  23. local COMPLETION_OPTIONS=""
  24. local BASH_VERSION_ARR=(${BASH_VERSION//./ })
  25. # Only BASH version 4.4 and later have the nosort option.
  26. if [ ${BASH_VERSION_ARR[0]} -gt 4 ] || ([ ${BASH_VERSION_ARR[0]} -eq 4 ] && [ ${BASH_VERSION_ARR[1]} -ge 4 ]); then
  27. COMPLETION_OPTIONS="-o nosort"
  28. fi
  29. complete $COMPLETION_OPTIONS -F %(complete_func)s %(script_names)s
  30. }
  31. %(complete_func)setup
  32. '''
  33. COMPLETION_SCRIPT_ZSH = '''
  34. %(complete_func)s() {
  35. local -a completions
  36. local -a completions_with_descriptions
  37. local -a response
  38. response=("${(@f)$( env COMP_WORDS=\"${words[*]}\" \\
  39. COMP_CWORD=$((CURRENT-1)) \\
  40. %(autocomplete_var)s=\"complete_zsh\" \\
  41. %(script_names)s )}")
  42. for key descr in ${(kv)response}; do
  43. if [[ "$descr" == "_" ]]; then
  44. completions+=("$key")
  45. else
  46. completions_with_descriptions+=("$key":"$descr")
  47. fi
  48. done
  49. if [ -n "$completions_with_descriptions" ]; then
  50. _describe -V unsorted completions_with_descriptions -U -Q
  51. fi
  52. if [ -n "$completions" ]; then
  53. compadd -U -V unsorted -Q -a completions
  54. fi
  55. compstate[insert]="automenu"
  56. }
  57. compdef %(complete_func)s %(script_names)s
  58. '''
  59. _invalid_ident_char_re = re.compile(r'[^a-zA-Z0-9_]')
  60. def get_completion_script(prog_name, complete_var, shell):
  61. cf_name = _invalid_ident_char_re.sub('', prog_name.replace('-', '_'))
  62. script = COMPLETION_SCRIPT_ZSH if shell == 'zsh' else COMPLETION_SCRIPT_BASH
  63. return (script % {
  64. 'complete_func': '_%s_completion' % cf_name,
  65. 'script_names': prog_name,
  66. 'autocomplete_var': complete_var,
  67. }).strip() + ';'
  68. def resolve_ctx(cli, prog_name, args):
  69. """
  70. Parse into a hierarchy of contexts. Contexts are connected through the parent variable.
  71. :param cli: command definition
  72. :param prog_name: the program that is running
  73. :param args: full list of args
  74. :return: the final context/command parsed
  75. """
  76. ctx = cli.make_context(prog_name, args, resilient_parsing=True)
  77. args = ctx.protected_args + ctx.args
  78. while args:
  79. if isinstance(ctx.command, MultiCommand):
  80. if not ctx.command.chain:
  81. cmd_name, cmd, args = ctx.command.resolve_command(ctx, args)
  82. if cmd is None:
  83. return ctx
  84. ctx = cmd.make_context(cmd_name, args, parent=ctx,
  85. resilient_parsing=True)
  86. args = ctx.protected_args + ctx.args
  87. else:
  88. # Walk chained subcommand contexts saving the last one.
  89. while args:
  90. cmd_name, cmd, args = ctx.command.resolve_command(ctx, args)
  91. if cmd is None:
  92. return ctx
  93. sub_ctx = cmd.make_context(cmd_name, args, parent=ctx,
  94. allow_extra_args=True,
  95. allow_interspersed_args=False,
  96. resilient_parsing=True)
  97. args = sub_ctx.args
  98. ctx = sub_ctx
  99. args = sub_ctx.protected_args + sub_ctx.args
  100. else:
  101. break
  102. return ctx
  103. def start_of_option(param_str):
  104. """
  105. :param param_str: param_str to check
  106. :return: whether or not this is the start of an option declaration (i.e. starts "-" or "--")
  107. """
  108. return param_str and param_str[:1] == '-'
  109. def is_incomplete_option(all_args, cmd_param):
  110. """
  111. :param all_args: the full original list of args supplied
  112. :param cmd_param: the current command paramter
  113. :return: whether or not the last option declaration (i.e. starts "-" or "--") is incomplete and
  114. corresponds to this cmd_param. In other words whether this cmd_param option can still accept
  115. values
  116. """
  117. if not isinstance(cmd_param, Option):
  118. return False
  119. if cmd_param.is_flag:
  120. return False
  121. last_option = None
  122. for index, arg_str in enumerate(reversed([arg for arg in all_args if arg != WORDBREAK])):
  123. if index + 1 > cmd_param.nargs:
  124. break
  125. if start_of_option(arg_str):
  126. last_option = arg_str
  127. return True if last_option and last_option in cmd_param.opts else False
  128. def is_incomplete_argument(current_params, cmd_param):
  129. """
  130. :param current_params: the current params and values for this argument as already entered
  131. :param cmd_param: the current command parameter
  132. :return: whether or not the last argument is incomplete and corresponds to this cmd_param. In
  133. other words whether or not the this cmd_param argument can still accept values
  134. """
  135. if not isinstance(cmd_param, Argument):
  136. return False
  137. current_param_values = current_params[cmd_param.name]
  138. if current_param_values is None:
  139. return True
  140. if cmd_param.nargs == -1:
  141. return True
  142. if isinstance(current_param_values, abc.Iterable) \
  143. and cmd_param.nargs > 1 and len(current_param_values) < cmd_param.nargs:
  144. return True
  145. return False
  146. def get_user_autocompletions(ctx, args, incomplete, cmd_param):
  147. """
  148. :param ctx: context associated with the parsed command
  149. :param args: full list of args
  150. :param incomplete: the incomplete text to autocomplete
  151. :param cmd_param: command definition
  152. :return: all the possible user-specified completions for the param
  153. """
  154. results = []
  155. if isinstance(cmd_param.type, Choice):
  156. # Choices don't support descriptions.
  157. results = [(c, None)
  158. for c in cmd_param.type.choices if str(c).startswith(incomplete)]
  159. elif cmd_param.autocompletion is not None:
  160. dynamic_completions = cmd_param.autocompletion(ctx=ctx,
  161. args=args,
  162. incomplete=incomplete)
  163. results = [c if isinstance(c, tuple) else (c, None)
  164. for c in dynamic_completions]
  165. return results
  166. def get_visible_commands_starting_with(ctx, starts_with):
  167. """
  168. :param ctx: context associated with the parsed command
  169. :starts_with: string that visible commands must start with.
  170. :return: all visible (not hidden) commands that start with starts_with.
  171. """
  172. for c in ctx.command.list_commands(ctx):
  173. if c.startswith(starts_with):
  174. command = ctx.command.get_command(ctx, c)
  175. if not command.hidden:
  176. yield command
  177. def add_subcommand_completions(ctx, incomplete, completions_out):
  178. # Add subcommand completions.
  179. if isinstance(ctx.command, MultiCommand):
  180. completions_out.extend(
  181. [(c.name, c.get_short_help_str()) for c in get_visible_commands_starting_with(ctx, incomplete)])
  182. # Walk up the context list and add any other completion possibilities from chained commands
  183. while ctx.parent is not None:
  184. ctx = ctx.parent
  185. if isinstance(ctx.command, MultiCommand) and ctx.command.chain:
  186. remaining_commands = [c for c in get_visible_commands_starting_with(ctx, incomplete)
  187. if c.name not in ctx.protected_args]
  188. completions_out.extend([(c.name, c.get_short_help_str()) for c in remaining_commands])
  189. def get_choices(cli, prog_name, args, incomplete):
  190. """
  191. :param cli: command definition
  192. :param prog_name: the program that is running
  193. :param args: full list of args
  194. :param incomplete: the incomplete text to autocomplete
  195. :return: all the possible completions for the incomplete
  196. """
  197. all_args = copy.deepcopy(args)
  198. ctx = resolve_ctx(cli, prog_name, args)
  199. if ctx is None:
  200. return []
  201. # In newer versions of bash long opts with '='s are partitioned, but it's easier to parse
  202. # without the '='
  203. if start_of_option(incomplete) and WORDBREAK in incomplete:
  204. partition_incomplete = incomplete.partition(WORDBREAK)
  205. all_args.append(partition_incomplete[0])
  206. incomplete = partition_incomplete[2]
  207. elif incomplete == WORDBREAK:
  208. incomplete = ''
  209. completions = []
  210. if start_of_option(incomplete):
  211. # completions for partial options
  212. for param in ctx.command.params:
  213. if isinstance(param, Option) and not param.hidden:
  214. param_opts = [param_opt for param_opt in param.opts +
  215. param.secondary_opts if param_opt not in all_args or param.multiple]
  216. completions.extend([(o, param.help) for o in param_opts if o.startswith(incomplete)])
  217. return completions
  218. # completion for option values from user supplied values
  219. for param in ctx.command.params:
  220. if is_incomplete_option(all_args, param):
  221. return get_user_autocompletions(ctx, all_args, incomplete, param)
  222. # completion for argument values from user supplied values
  223. for param in ctx.command.params:
  224. if is_incomplete_argument(ctx.params, param):
  225. return get_user_autocompletions(ctx, all_args, incomplete, param)
  226. add_subcommand_completions(ctx, incomplete, completions)
  227. # Sort before returning so that proper ordering can be enforced in custom types.
  228. return sorted(completions)
  229. def do_complete(cli, prog_name, include_descriptions):
  230. cwords = split_arg_string(os.environ['COMP_WORDS'])
  231. cword = int(os.environ['COMP_CWORD'])
  232. args = cwords[1:cword]
  233. try:
  234. incomplete = cwords[cword]
  235. except IndexError:
  236. incomplete = ''
  237. for item in get_choices(cli, prog_name, args, incomplete):
  238. echo(item[0])
  239. if include_descriptions:
  240. # ZSH has trouble dealing with empty array parameters when returned from commands, so use a well defined character '_' to indicate no description is present.
  241. echo(item[1] if item[1] else '_')
  242. return True
  243. def bashcomplete(cli, prog_name, complete_var, complete_instr):
  244. if complete_instr.startswith('source'):
  245. shell = 'zsh' if complete_instr == 'source_zsh' else 'bash'
  246. echo(get_completion_script(prog_name, complete_var, shell))
  247. return True
  248. elif complete_instr == 'complete' or complete_instr == 'complete_zsh':
  249. return do_complete(cli, prog_name, complete_instr == 'complete_zsh')
  250. return False