shell_completion.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596
  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.MutableMapping[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: t.Any = value
  69. self.type: str = type
  70. self.help: t.Optional[str] = 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. if [[ $zsh_eval_context[-1] == loadautofunc ]]; then
  130. # autoload from fpath, call function directly
  131. %(complete_func)s "$@"
  132. else
  133. # eval/source/. command, register function for later
  134. compdef %(complete_func)s %(prog_name)s
  135. fi
  136. """
  137. _SOURCE_FISH = """\
  138. function %(complete_func)s;
  139. set -l response (env %(complete_var)s=fish_complete COMP_WORDS=(commandline -cp) \
  140. COMP_CWORD=(commandline -t) %(prog_name)s);
  141. for completion in $response;
  142. set -l metadata (string split "," $completion);
  143. if test $metadata[1] = "dir";
  144. __fish_complete_directories $metadata[2];
  145. else if test $metadata[1] = "file";
  146. __fish_complete_path $metadata[2];
  147. else if test $metadata[1] = "plain";
  148. echo $metadata[2];
  149. end;
  150. end;
  151. end;
  152. complete --no-files --command %(prog_name)s --arguments \
  153. "(%(complete_func)s)";
  154. """
  155. class ShellComplete:
  156. """Base class for providing shell completion support. A subclass for
  157. a given shell will override attributes and methods to implement the
  158. completion instructions (``source`` and ``complete``).
  159. :param cli: Command being called.
  160. :param prog_name: Name of the executable in the shell.
  161. :param complete_var: Name of the environment variable that holds
  162. the completion instruction.
  163. .. versionadded:: 8.0
  164. """
  165. name: t.ClassVar[str]
  166. """Name to register the shell as with :func:`add_completion_class`.
  167. This is used in completion instructions (``{name}_source`` and
  168. ``{name}_complete``).
  169. """
  170. source_template: t.ClassVar[str]
  171. """Completion script template formatted by :meth:`source`. This must
  172. be provided by subclasses.
  173. """
  174. def __init__(
  175. self,
  176. cli: BaseCommand,
  177. ctx_args: t.MutableMapping[str, t.Any],
  178. prog_name: str,
  179. complete_var: str,
  180. ) -> None:
  181. self.cli = cli
  182. self.ctx_args = ctx_args
  183. self.prog_name = prog_name
  184. self.complete_var = complete_var
  185. @property
  186. def func_name(self) -> str:
  187. """The name of the shell function defined by the completion
  188. script.
  189. """
  190. safe_name = re.sub(r"\W*", "", self.prog_name.replace("-", "_"), flags=re.ASCII)
  191. return f"_{safe_name}_completion"
  192. def source_vars(self) -> t.Dict[str, t.Any]:
  193. """Vars for formatting :attr:`source_template`.
  194. By default this provides ``complete_func``, ``complete_var``,
  195. and ``prog_name``.
  196. """
  197. return {
  198. "complete_func": self.func_name,
  199. "complete_var": self.complete_var,
  200. "prog_name": self.prog_name,
  201. }
  202. def source(self) -> str:
  203. """Produce the shell script that defines the completion
  204. function. By default this ``%``-style formats
  205. :attr:`source_template` with the dict returned by
  206. :meth:`source_vars`.
  207. """
  208. return self.source_template % self.source_vars()
  209. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  210. """Use the env vars defined by the shell script to return a
  211. tuple of ``args, incomplete``. This must be implemented by
  212. subclasses.
  213. """
  214. raise NotImplementedError
  215. def get_completions(
  216. self, args: t.List[str], incomplete: str
  217. ) -> t.List[CompletionItem]:
  218. """Determine the context and last complete command or parameter
  219. from the complete args. Call that object's ``shell_complete``
  220. method to get the completions for the incomplete value.
  221. :param args: List of complete args before the incomplete value.
  222. :param incomplete: Value being completed. May be empty.
  223. """
  224. ctx = _resolve_context(self.cli, self.ctx_args, self.prog_name, args)
  225. obj, incomplete = _resolve_incomplete(ctx, args, incomplete)
  226. return obj.shell_complete(ctx, incomplete)
  227. def format_completion(self, item: CompletionItem) -> str:
  228. """Format a completion item into the form recognized by the
  229. shell script. This must be implemented by subclasses.
  230. :param item: Completion item to format.
  231. """
  232. raise NotImplementedError
  233. def complete(self) -> str:
  234. """Produce the completion data to send back to the shell.
  235. By default this calls :meth:`get_completion_args`, gets the
  236. completions, then calls :meth:`format_completion` for each
  237. completion.
  238. """
  239. args, incomplete = self.get_completion_args()
  240. completions = self.get_completions(args, incomplete)
  241. out = [self.format_completion(item) for item in completions]
  242. return "\n".join(out)
  243. class BashComplete(ShellComplete):
  244. """Shell completion for Bash."""
  245. name = "bash"
  246. source_template = _SOURCE_BASH
  247. @staticmethod
  248. def _check_version() -> None:
  249. import subprocess
  250. output = subprocess.run(
  251. ["bash", "-c", 'echo "${BASH_VERSION}"'], stdout=subprocess.PIPE
  252. )
  253. match = re.search(r"^(\d+)\.(\d+)\.\d+", output.stdout.decode())
  254. if match is not None:
  255. major, minor = match.groups()
  256. if major < "4" or major == "4" and minor < "4":
  257. echo(
  258. _(
  259. "Shell completion is not supported for Bash"
  260. " versions older than 4.4."
  261. ),
  262. err=True,
  263. )
  264. else:
  265. echo(
  266. _("Couldn't detect Bash version, shell completion is not supported."),
  267. err=True,
  268. )
  269. def source(self) -> str:
  270. self._check_version()
  271. return super().source()
  272. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  273. cwords = split_arg_string(os.environ["COMP_WORDS"])
  274. cword = int(os.environ["COMP_CWORD"])
  275. args = cwords[1:cword]
  276. try:
  277. incomplete = cwords[cword]
  278. except IndexError:
  279. incomplete = ""
  280. return args, incomplete
  281. def format_completion(self, item: CompletionItem) -> str:
  282. return f"{item.type},{item.value}"
  283. class ZshComplete(ShellComplete):
  284. """Shell completion for Zsh."""
  285. name = "zsh"
  286. source_template = _SOURCE_ZSH
  287. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  288. cwords = split_arg_string(os.environ["COMP_WORDS"])
  289. cword = int(os.environ["COMP_CWORD"])
  290. args = cwords[1:cword]
  291. try:
  292. incomplete = cwords[cword]
  293. except IndexError:
  294. incomplete = ""
  295. return args, incomplete
  296. def format_completion(self, item: CompletionItem) -> str:
  297. return f"{item.type}\n{item.value}\n{item.help if item.help else '_'}"
  298. class FishComplete(ShellComplete):
  299. """Shell completion for Fish."""
  300. name = "fish"
  301. source_template = _SOURCE_FISH
  302. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  303. cwords = split_arg_string(os.environ["COMP_WORDS"])
  304. incomplete = os.environ["COMP_CWORD"]
  305. args = cwords[1:]
  306. # Fish stores the partial word in both COMP_WORDS and
  307. # COMP_CWORD, remove it from complete args.
  308. if incomplete and args and args[-1] == incomplete:
  309. args.pop()
  310. return args, incomplete
  311. def format_completion(self, item: CompletionItem) -> str:
  312. if item.help:
  313. return f"{item.type},{item.value}\t{item.help}"
  314. return f"{item.type},{item.value}"
  315. ShellCompleteType = t.TypeVar("ShellCompleteType", bound=t.Type[ShellComplete])
  316. _available_shells: t.Dict[str, t.Type[ShellComplete]] = {
  317. "bash": BashComplete,
  318. "fish": FishComplete,
  319. "zsh": ZshComplete,
  320. }
  321. def add_completion_class(
  322. cls: ShellCompleteType, name: t.Optional[str] = None
  323. ) -> ShellCompleteType:
  324. """Register a :class:`ShellComplete` subclass under the given name.
  325. The name will be provided by the completion instruction environment
  326. variable during completion.
  327. :param cls: The completion class that will handle completion for the
  328. shell.
  329. :param name: Name to register the class under. Defaults to the
  330. class's ``name`` attribute.
  331. """
  332. if name is None:
  333. name = cls.name
  334. _available_shells[name] = cls
  335. return cls
  336. def get_completion_class(shell: str) -> t.Optional[t.Type[ShellComplete]]:
  337. """Look up a registered :class:`ShellComplete` subclass by the name
  338. provided by the completion instruction environment variable. If the
  339. name isn't registered, returns ``None``.
  340. :param shell: Name the class is registered under.
  341. """
  342. return _available_shells.get(shell)
  343. def _is_incomplete_argument(ctx: Context, param: Parameter) -> bool:
  344. """Determine if the given parameter is an argument that can still
  345. accept values.
  346. :param ctx: Invocation context for the command represented by the
  347. parsed complete args.
  348. :param param: Argument object being checked.
  349. """
  350. if not isinstance(param, Argument):
  351. return False
  352. assert param.name is not None
  353. # Will be None if expose_value is False.
  354. value = ctx.params.get(param.name)
  355. return (
  356. param.nargs == -1
  357. or ctx.get_parameter_source(param.name) is not ParameterSource.COMMANDLINE
  358. or (
  359. param.nargs > 1
  360. and isinstance(value, (tuple, list))
  361. and len(value) < param.nargs
  362. )
  363. )
  364. def _start_of_option(ctx: Context, value: str) -> bool:
  365. """Check if the value looks like the start of an option."""
  366. if not value:
  367. return False
  368. c = value[0]
  369. return c in ctx._opt_prefixes
  370. def _is_incomplete_option(ctx: Context, args: t.List[str], param: Parameter) -> bool:
  371. """Determine if the given parameter is an option that needs a value.
  372. :param args: List of complete args before the incomplete value.
  373. :param param: Option object being checked.
  374. """
  375. if not isinstance(param, Option):
  376. return False
  377. if param.is_flag or param.count:
  378. return False
  379. last_option = None
  380. for index, arg in enumerate(reversed(args)):
  381. if index + 1 > param.nargs:
  382. break
  383. if _start_of_option(ctx, arg):
  384. last_option = arg
  385. return last_option is not None and last_option in param.opts
  386. def _resolve_context(
  387. cli: BaseCommand,
  388. ctx_args: t.MutableMapping[str, t.Any],
  389. prog_name: str,
  390. args: t.List[str],
  391. ) -> Context:
  392. """Produce the context hierarchy starting with the command and
  393. traversing the complete arguments. This only follows the commands,
  394. it doesn't trigger input prompts or callbacks.
  395. :param cli: Command being called.
  396. :param prog_name: Name of the executable in the shell.
  397. :param args: List of complete args before the incomplete value.
  398. """
  399. ctx_args["resilient_parsing"] = True
  400. ctx = cli.make_context(prog_name, args.copy(), **ctx_args)
  401. args = ctx.protected_args + ctx.args
  402. while args:
  403. command = ctx.command
  404. if isinstance(command, MultiCommand):
  405. if not command.chain:
  406. name, cmd, args = command.resolve_command(ctx, args)
  407. if cmd is None:
  408. return ctx
  409. ctx = cmd.make_context(name, args, parent=ctx, resilient_parsing=True)
  410. args = ctx.protected_args + ctx.args
  411. else:
  412. sub_ctx = ctx
  413. while args:
  414. name, cmd, args = command.resolve_command(ctx, args)
  415. if cmd is None:
  416. return ctx
  417. sub_ctx = cmd.make_context(
  418. name,
  419. args,
  420. parent=ctx,
  421. allow_extra_args=True,
  422. allow_interspersed_args=False,
  423. resilient_parsing=True,
  424. )
  425. args = sub_ctx.args
  426. ctx = sub_ctx
  427. args = [*sub_ctx.protected_args, *sub_ctx.args]
  428. else:
  429. break
  430. return ctx
  431. def _resolve_incomplete(
  432. ctx: Context, args: t.List[str], incomplete: str
  433. ) -> t.Tuple[t.Union[BaseCommand, Parameter], str]:
  434. """Find the Click object that will handle the completion of the
  435. incomplete value. Return the object and the incomplete value.
  436. :param ctx: Invocation context for the command represented by
  437. the parsed complete args.
  438. :param args: List of complete args before the incomplete value.
  439. :param incomplete: Value being completed. May be empty.
  440. """
  441. # Different shells treat an "=" between a long option name and
  442. # value differently. Might keep the value joined, return the "="
  443. # as a separate item, or return the split name and value. Always
  444. # split and discard the "=" to make completion easier.
  445. if incomplete == "=":
  446. incomplete = ""
  447. elif "=" in incomplete and _start_of_option(ctx, incomplete):
  448. name, _, incomplete = incomplete.partition("=")
  449. args.append(name)
  450. # The "--" marker tells Click to stop treating values as options
  451. # even if they start with the option character. If it hasn't been
  452. # given and the incomplete arg looks like an option, the current
  453. # command will provide option name completions.
  454. if "--" not in args and _start_of_option(ctx, incomplete):
  455. return ctx.command, incomplete
  456. params = ctx.command.get_params(ctx)
  457. # If the last complete arg is an option name with an incomplete
  458. # value, the option will provide value completions.
  459. for param in params:
  460. if _is_incomplete_option(ctx, args, param):
  461. return param, incomplete
  462. # It's not an option name or value. The first argument without a
  463. # parsed value will provide value completions.
  464. for param in params:
  465. if _is_incomplete_argument(ctx, param):
  466. return param, incomplete
  467. # There were no unparsed arguments, the command may be a group that
  468. # will provide command name completions.
  469. return ctx.command, incomplete