shell_completion.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580
  1. import os
  2. import re
  3. import typing as t
  4. from gettext import gettext as _
  5. from .core import Argument
  6. from .core import BaseCommand
  7. from .core import Context
  8. from .core import MultiCommand
  9. from .core import Option
  10. from .core import Parameter
  11. from .core import ParameterSource
  12. from .parser import split_arg_string
  13. from .utils import echo
  14. def shell_complete(
  15. cli: BaseCommand,
  16. ctx_args: t.Dict[str, t.Any],
  17. prog_name: str,
  18. complete_var: str,
  19. instruction: str,
  20. ) -> int:
  21. """Perform shell completion for the given CLI program.
  22. :param cli: Command being called.
  23. :param ctx_args: Extra arguments to pass to
  24. ``cli.make_context``.
  25. :param prog_name: Name of the executable in the shell.
  26. :param complete_var: Name of the environment variable that holds
  27. the completion instruction.
  28. :param instruction: Value of ``complete_var`` with the completion
  29. instruction and shell, in the form ``instruction_shell``.
  30. :return: Status code to exit with.
  31. """
  32. shell, _, instruction = instruction.partition("_")
  33. comp_cls = get_completion_class(shell)
  34. if comp_cls is None:
  35. return 1
  36. comp = comp_cls(cli, ctx_args, prog_name, complete_var)
  37. if instruction == "source":
  38. echo(comp.source())
  39. return 0
  40. if instruction == "complete":
  41. echo(comp.complete())
  42. return 0
  43. return 1
  44. class CompletionItem:
  45. """Represents a completion value and metadata about the value. The
  46. default metadata is ``type`` to indicate special shell handling,
  47. and ``help`` if a shell supports showing a help string next to the
  48. value.
  49. Arbitrary parameters can be passed when creating the object, and
  50. accessed using ``item.attr``. If an attribute wasn't passed,
  51. accessing it returns ``None``.
  52. :param value: The completion suggestion.
  53. :param type: Tells the shell script to provide special completion
  54. support for the type. Click uses ``"dir"`` and ``"file"``.
  55. :param help: String shown next to the value if supported.
  56. :param kwargs: Arbitrary metadata. The built-in implementations
  57. don't use this, but custom type completions paired with custom
  58. shell support could use it.
  59. """
  60. __slots__ = ("value", "type", "help", "_info")
  61. def __init__(
  62. self,
  63. value: t.Any,
  64. type: str = "plain",
  65. help: t.Optional[str] = None,
  66. **kwargs: t.Any,
  67. ) -> None:
  68. self.value = value
  69. self.type = type
  70. self.help = help
  71. self._info = kwargs
  72. def __getattr__(self, name: str) -> t.Any:
  73. return self._info.get(name)
  74. # Only Bash >= 4.4 has the nosort option.
  75. _SOURCE_BASH = """\
  76. %(complete_func)s() {
  77. local IFS=$'\\n'
  78. local response
  79. response=$(env COMP_WORDS="${COMP_WORDS[*]}" COMP_CWORD=$COMP_CWORD \
  80. %(complete_var)s=bash_complete $1)
  81. for completion in $response; do
  82. IFS=',' read type value <<< "$completion"
  83. if [[ $type == 'dir' ]]; then
  84. COMPREPLY=()
  85. compopt -o dirnames
  86. elif [[ $type == 'file' ]]; then
  87. COMPREPLY=()
  88. compopt -o default
  89. elif [[ $type == 'plain' ]]; then
  90. COMPREPLY+=($value)
  91. fi
  92. done
  93. return 0
  94. }
  95. %(complete_func)s_setup() {
  96. complete -o nosort -F %(complete_func)s %(prog_name)s
  97. }
  98. %(complete_func)s_setup;
  99. """
  100. _SOURCE_ZSH = """\
  101. #compdef %(prog_name)s
  102. %(complete_func)s() {
  103. local -a completions
  104. local -a completions_with_descriptions
  105. local -a response
  106. (( ! $+commands[%(prog_name)s] )) && return 1
  107. response=("${(@f)$(env COMP_WORDS="${words[*]}" COMP_CWORD=$((CURRENT-1)) \
  108. %(complete_var)s=zsh_complete %(prog_name)s)}")
  109. for type key descr in ${response}; do
  110. if [[ "$type" == "plain" ]]; then
  111. if [[ "$descr" == "_" ]]; then
  112. completions+=("$key")
  113. else
  114. completions_with_descriptions+=("$key":"$descr")
  115. fi
  116. elif [[ "$type" == "dir" ]]; then
  117. _path_files -/
  118. elif [[ "$type" == "file" ]]; then
  119. _path_files -f
  120. fi
  121. done
  122. if [ -n "$completions_with_descriptions" ]; then
  123. _describe -V unsorted completions_with_descriptions -U
  124. fi
  125. if [ -n "$completions" ]; then
  126. compadd -U -V unsorted -a completions
  127. fi
  128. }
  129. compdef %(complete_func)s %(prog_name)s;
  130. """
  131. _SOURCE_FISH = """\
  132. function %(complete_func)s;
  133. set -l response;
  134. for value in (env %(complete_var)s=fish_complete COMP_WORDS=(commandline -cp) \
  135. COMP_CWORD=(commandline -t) %(prog_name)s);
  136. set response $response $value;
  137. end;
  138. for completion in $response;
  139. set -l metadata (string split "," $completion);
  140. if test $metadata[1] = "dir";
  141. __fish_complete_directories $metadata[2];
  142. else if test $metadata[1] = "file";
  143. __fish_complete_path $metadata[2];
  144. else if test $metadata[1] = "plain";
  145. echo $metadata[2];
  146. end;
  147. end;
  148. end;
  149. complete --no-files --command %(prog_name)s --arguments \
  150. "(%(complete_func)s)";
  151. """
  152. class ShellComplete:
  153. """Base class for providing shell completion support. A subclass for
  154. a given shell will override attributes and methods to implement the
  155. completion instructions (``source`` and ``complete``).
  156. :param cli: Command being called.
  157. :param prog_name: Name of the executable in the shell.
  158. :param complete_var: Name of the environment variable that holds
  159. the completion instruction.
  160. .. versionadded:: 8.0
  161. """
  162. name: t.ClassVar[str]
  163. """Name to register the shell as with :func:`add_completion_class`.
  164. This is used in completion instructions (``{name}_source`` and
  165. ``{name}_complete``).
  166. """
  167. source_template: t.ClassVar[str]
  168. """Completion script template formatted by :meth:`source`. This must
  169. be provided by subclasses.
  170. """
  171. def __init__(
  172. self,
  173. cli: BaseCommand,
  174. ctx_args: t.Dict[str, t.Any],
  175. prog_name: str,
  176. complete_var: str,
  177. ) -> None:
  178. self.cli = cli
  179. self.ctx_args = ctx_args
  180. self.prog_name = prog_name
  181. self.complete_var = complete_var
  182. @property
  183. def func_name(self) -> str:
  184. """The name of the shell function defined by the completion
  185. script.
  186. """
  187. safe_name = re.sub(r"\W*", "", self.prog_name.replace("-", "_"), re.ASCII)
  188. return f"_{safe_name}_completion"
  189. def source_vars(self) -> t.Dict[str, t.Any]:
  190. """Vars for formatting :attr:`source_template`.
  191. By default this provides ``complete_func``, ``complete_var``,
  192. and ``prog_name``.
  193. """
  194. return {
  195. "complete_func": self.func_name,
  196. "complete_var": self.complete_var,
  197. "prog_name": self.prog_name,
  198. }
  199. def source(self) -> str:
  200. """Produce the shell script that defines the completion
  201. function. By default this ``%``-style formats
  202. :attr:`source_template` with the dict returned by
  203. :meth:`source_vars`.
  204. """
  205. return self.source_template % self.source_vars()
  206. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  207. """Use the env vars defined by the shell script to return a
  208. tuple of ``args, incomplete``. This must be implemented by
  209. subclasses.
  210. """
  211. raise NotImplementedError
  212. def get_completions(
  213. self, args: t.List[str], incomplete: str
  214. ) -> t.List[CompletionItem]:
  215. """Determine the context and last complete command or parameter
  216. from the complete args. Call that object's ``shell_complete``
  217. method to get the completions for the incomplete value.
  218. :param args: List of complete args before the incomplete value.
  219. :param incomplete: Value being completed. May be empty.
  220. """
  221. ctx = _resolve_context(self.cli, self.ctx_args, self.prog_name, args)
  222. obj, incomplete = _resolve_incomplete(ctx, args, incomplete)
  223. return obj.shell_complete(ctx, incomplete)
  224. def format_completion(self, item: CompletionItem) -> str:
  225. """Format a completion item into the form recognized by the
  226. shell script. This must be implemented by subclasses.
  227. :param item: Completion item to format.
  228. """
  229. raise NotImplementedError
  230. def complete(self) -> str:
  231. """Produce the completion data to send back to the shell.
  232. By default this calls :meth:`get_completion_args`, gets the
  233. completions, then calls :meth:`format_completion` for each
  234. completion.
  235. """
  236. args, incomplete = self.get_completion_args()
  237. completions = self.get_completions(args, incomplete)
  238. out = [self.format_completion(item) for item in completions]
  239. return "\n".join(out)
  240. class BashComplete(ShellComplete):
  241. """Shell completion for Bash."""
  242. name = "bash"
  243. source_template = _SOURCE_BASH
  244. def _check_version(self) -> None:
  245. import subprocess
  246. output = subprocess.run(
  247. ["bash", "-c", "echo ${BASH_VERSION}"], stdout=subprocess.PIPE
  248. )
  249. match = re.search(r"^(\d+)\.(\d+)\.\d+", output.stdout.decode())
  250. if match is not None:
  251. major, minor = match.groups()
  252. if major < "4" or major == "4" and minor < "4":
  253. raise RuntimeError(
  254. _(
  255. "Shell completion is not supported for Bash"
  256. " versions older than 4.4."
  257. )
  258. )
  259. else:
  260. raise RuntimeError(
  261. _("Couldn't detect Bash version, shell completion is not supported.")
  262. )
  263. def source(self) -> str:
  264. self._check_version()
  265. return super().source()
  266. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  267. cwords = split_arg_string(os.environ["COMP_WORDS"])
  268. cword = int(os.environ["COMP_CWORD"])
  269. args = cwords[1:cword]
  270. try:
  271. incomplete = cwords[cword]
  272. except IndexError:
  273. incomplete = ""
  274. return args, incomplete
  275. def format_completion(self, item: CompletionItem) -> str:
  276. return f"{item.type},{item.value}"
  277. class ZshComplete(ShellComplete):
  278. """Shell completion for Zsh."""
  279. name = "zsh"
  280. source_template = _SOURCE_ZSH
  281. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  282. cwords = split_arg_string(os.environ["COMP_WORDS"])
  283. cword = int(os.environ["COMP_CWORD"])
  284. args = cwords[1:cword]
  285. try:
  286. incomplete = cwords[cword]
  287. except IndexError:
  288. incomplete = ""
  289. return args, incomplete
  290. def format_completion(self, item: CompletionItem) -> str:
  291. return f"{item.type}\n{item.value}\n{item.help if item.help else '_'}"
  292. class FishComplete(ShellComplete):
  293. """Shell completion for Fish."""
  294. name = "fish"
  295. source_template = _SOURCE_FISH
  296. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  297. cwords = split_arg_string(os.environ["COMP_WORDS"])
  298. incomplete = os.environ["COMP_CWORD"]
  299. args = cwords[1:]
  300. # Fish stores the partial word in both COMP_WORDS and
  301. # COMP_CWORD, remove it from complete args.
  302. if incomplete and args and args[-1] == incomplete:
  303. args.pop()
  304. return args, incomplete
  305. def format_completion(self, item: CompletionItem) -> str:
  306. if item.help:
  307. return f"{item.type},{item.value}\t{item.help}"
  308. return f"{item.type},{item.value}"
  309. _available_shells: t.Dict[str, t.Type[ShellComplete]] = {
  310. "bash": BashComplete,
  311. "fish": FishComplete,
  312. "zsh": ZshComplete,
  313. }
  314. def add_completion_class(
  315. cls: t.Type[ShellComplete], name: t.Optional[str] = None
  316. ) -> None:
  317. """Register a :class:`ShellComplete` subclass under the given name.
  318. The name will be provided by the completion instruction environment
  319. variable during completion.
  320. :param cls: The completion class that will handle completion for the
  321. shell.
  322. :param name: Name to register the class under. Defaults to the
  323. class's ``name`` attribute.
  324. """
  325. if name is None:
  326. name = cls.name
  327. _available_shells[name] = cls
  328. def get_completion_class(shell: str) -> t.Optional[t.Type[ShellComplete]]:
  329. """Look up a registered :class:`ShellComplete` subclass by the name
  330. provided by the completion instruction environment variable. If the
  331. name isn't registered, returns ``None``.
  332. :param shell: Name the class is registered under.
  333. """
  334. return _available_shells.get(shell)
  335. def _is_incomplete_argument(ctx: Context, param: Parameter) -> bool:
  336. """Determine if the given parameter is an argument that can still
  337. accept values.
  338. :param ctx: Invocation context for the command represented by the
  339. parsed complete args.
  340. :param param: Argument object being checked.
  341. """
  342. if not isinstance(param, Argument):
  343. return False
  344. assert param.name is not None
  345. value = ctx.params[param.name]
  346. return (
  347. param.nargs == -1
  348. or ctx.get_parameter_source(param.name) is not ParameterSource.COMMANDLINE
  349. or (
  350. param.nargs > 1
  351. and isinstance(value, (tuple, list))
  352. and len(value) < param.nargs
  353. )
  354. )
  355. def _start_of_option(ctx: Context, value: str) -> bool:
  356. """Check if the value looks like the start of an option."""
  357. if not value:
  358. return False
  359. c = value[0]
  360. return c in ctx._opt_prefixes
  361. def _is_incomplete_option(ctx: Context, args: t.List[str], param: Parameter) -> bool:
  362. """Determine if the given parameter is an option that needs a value.
  363. :param args: List of complete args before the incomplete value.
  364. :param param: Option object being checked.
  365. """
  366. if not isinstance(param, Option):
  367. return False
  368. if param.is_flag or param.count:
  369. return False
  370. last_option = None
  371. for index, arg in enumerate(reversed(args)):
  372. if index + 1 > param.nargs:
  373. break
  374. if _start_of_option(ctx, arg):
  375. last_option = arg
  376. return last_option is not None and last_option in param.opts
  377. def _resolve_context(
  378. cli: BaseCommand, ctx_args: t.Dict[str, t.Any], prog_name: str, args: t.List[str]
  379. ) -> Context:
  380. """Produce the context hierarchy starting with the command and
  381. traversing the complete arguments. This only follows the commands,
  382. it doesn't trigger input prompts or callbacks.
  383. :param cli: Command being called.
  384. :param prog_name: Name of the executable in the shell.
  385. :param args: List of complete args before the incomplete value.
  386. """
  387. ctx_args["resilient_parsing"] = True
  388. ctx = cli.make_context(prog_name, args.copy(), **ctx_args)
  389. args = ctx.protected_args + ctx.args
  390. while args:
  391. command = ctx.command
  392. if isinstance(command, MultiCommand):
  393. if not command.chain:
  394. name, cmd, args = command.resolve_command(ctx, args)
  395. if cmd is None:
  396. return ctx
  397. ctx = cmd.make_context(name, args, parent=ctx, resilient_parsing=True)
  398. args = ctx.protected_args + ctx.args
  399. else:
  400. while args:
  401. name, cmd, args = command.resolve_command(ctx, args)
  402. if cmd is None:
  403. return ctx
  404. sub_ctx = cmd.make_context(
  405. name,
  406. args,
  407. parent=ctx,
  408. allow_extra_args=True,
  409. allow_interspersed_args=False,
  410. resilient_parsing=True,
  411. )
  412. args = sub_ctx.args
  413. ctx = sub_ctx
  414. args = [*sub_ctx.protected_args, *sub_ctx.args]
  415. else:
  416. break
  417. return ctx
  418. def _resolve_incomplete(
  419. ctx: Context, args: t.List[str], incomplete: str
  420. ) -> t.Tuple[t.Union[BaseCommand, Parameter], str]:
  421. """Find the Click object that will handle the completion of the
  422. incomplete value. Return the object and the incomplete value.
  423. :param ctx: Invocation context for the command represented by
  424. the parsed complete args.
  425. :param args: List of complete args before the incomplete value.
  426. :param incomplete: Value being completed. May be empty.
  427. """
  428. # Different shells treat an "=" between a long option name and
  429. # value differently. Might keep the value joined, return the "="
  430. # as a separate item, or return the split name and value. Always
  431. # split and discard the "=" to make completion easier.
  432. if incomplete == "=":
  433. incomplete = ""
  434. elif "=" in incomplete and _start_of_option(ctx, incomplete):
  435. name, _, incomplete = incomplete.partition("=")
  436. args.append(name)
  437. # The "--" marker tells Click to stop treating values as options
  438. # even if they start with the option character. If it hasn't been
  439. # given and the incomplete arg looks like an option, the current
  440. # command will provide option name completions.
  441. if "--" not in args and _start_of_option(ctx, incomplete):
  442. return ctx.command, incomplete
  443. params = ctx.command.get_params(ctx)
  444. # If the last complete arg is an option name with an incomplete
  445. # value, the option will provide value completions.
  446. for param in params:
  447. if _is_incomplete_option(ctx, args, param):
  448. return param, incomplete
  449. # It's not an option name or value. The first argument without a
  450. # parsed value will provide value completions.
  451. for param in params:
  452. if _is_incomplete_argument(ctx, param):
  453. return param, incomplete
  454. # There were no unparsed arguments, the command may be a group that
  455. # will provide command name completions.
  456. return ctx.command, incomplete