shell_completion.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603
  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 shutil
  250. import subprocess
  251. bash_exe = shutil.which("bash")
  252. if bash_exe is None:
  253. match = None
  254. else:
  255. output = subprocess.run(
  256. [bash_exe, "--norc", "-c", 'echo "${BASH_VERSION}"'],
  257. stdout=subprocess.PIPE,
  258. )
  259. match = re.search(r"^(\d+)\.(\d+)\.\d+", output.stdout.decode())
  260. if match is not None:
  261. major, minor = match.groups()
  262. if major < "4" or major == "4" and minor < "4":
  263. echo(
  264. _(
  265. "Shell completion is not supported for Bash"
  266. " versions older than 4.4."
  267. ),
  268. err=True,
  269. )
  270. else:
  271. echo(
  272. _("Couldn't detect Bash version, shell completion is not supported."),
  273. err=True,
  274. )
  275. def source(self) -> str:
  276. self._check_version()
  277. return super().source()
  278. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  279. cwords = split_arg_string(os.environ["COMP_WORDS"])
  280. cword = int(os.environ["COMP_CWORD"])
  281. args = cwords[1:cword]
  282. try:
  283. incomplete = cwords[cword]
  284. except IndexError:
  285. incomplete = ""
  286. return args, incomplete
  287. def format_completion(self, item: CompletionItem) -> str:
  288. return f"{item.type},{item.value}"
  289. class ZshComplete(ShellComplete):
  290. """Shell completion for Zsh."""
  291. name = "zsh"
  292. source_template = _SOURCE_ZSH
  293. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  294. cwords = split_arg_string(os.environ["COMP_WORDS"])
  295. cword = int(os.environ["COMP_CWORD"])
  296. args = cwords[1:cword]
  297. try:
  298. incomplete = cwords[cword]
  299. except IndexError:
  300. incomplete = ""
  301. return args, incomplete
  302. def format_completion(self, item: CompletionItem) -> str:
  303. return f"{item.type}\n{item.value}\n{item.help if item.help else '_'}"
  304. class FishComplete(ShellComplete):
  305. """Shell completion for Fish."""
  306. name = "fish"
  307. source_template = _SOURCE_FISH
  308. def get_completion_args(self) -> t.Tuple[t.List[str], str]:
  309. cwords = split_arg_string(os.environ["COMP_WORDS"])
  310. incomplete = os.environ["COMP_CWORD"]
  311. args = cwords[1:]
  312. # Fish stores the partial word in both COMP_WORDS and
  313. # COMP_CWORD, remove it from complete args.
  314. if incomplete and args and args[-1] == incomplete:
  315. args.pop()
  316. return args, incomplete
  317. def format_completion(self, item: CompletionItem) -> str:
  318. if item.help:
  319. return f"{item.type},{item.value}\t{item.help}"
  320. return f"{item.type},{item.value}"
  321. ShellCompleteType = t.TypeVar("ShellCompleteType", bound=t.Type[ShellComplete])
  322. _available_shells: t.Dict[str, t.Type[ShellComplete]] = {
  323. "bash": BashComplete,
  324. "fish": FishComplete,
  325. "zsh": ZshComplete,
  326. }
  327. def add_completion_class(
  328. cls: ShellCompleteType, name: t.Optional[str] = None
  329. ) -> ShellCompleteType:
  330. """Register a :class:`ShellComplete` subclass under the given name.
  331. The name will be provided by the completion instruction environment
  332. variable during completion.
  333. :param cls: The completion class that will handle completion for the
  334. shell.
  335. :param name: Name to register the class under. Defaults to the
  336. class's ``name`` attribute.
  337. """
  338. if name is None:
  339. name = cls.name
  340. _available_shells[name] = cls
  341. return cls
  342. def get_completion_class(shell: str) -> t.Optional[t.Type[ShellComplete]]:
  343. """Look up a registered :class:`ShellComplete` subclass by the name
  344. provided by the completion instruction environment variable. If the
  345. name isn't registered, returns ``None``.
  346. :param shell: Name the class is registered under.
  347. """
  348. return _available_shells.get(shell)
  349. def _is_incomplete_argument(ctx: Context, param: Parameter) -> bool:
  350. """Determine if the given parameter is an argument that can still
  351. accept values.
  352. :param ctx: Invocation context for the command represented by the
  353. parsed complete args.
  354. :param param: Argument object being checked.
  355. """
  356. if not isinstance(param, Argument):
  357. return False
  358. assert param.name is not None
  359. # Will be None if expose_value is False.
  360. value = ctx.params.get(param.name)
  361. return (
  362. param.nargs == -1
  363. or ctx.get_parameter_source(param.name) is not ParameterSource.COMMANDLINE
  364. or (
  365. param.nargs > 1
  366. and isinstance(value, (tuple, list))
  367. and len(value) < param.nargs
  368. )
  369. )
  370. def _start_of_option(ctx: Context, value: str) -> bool:
  371. """Check if the value looks like the start of an option."""
  372. if not value:
  373. return False
  374. c = value[0]
  375. return c in ctx._opt_prefixes
  376. def _is_incomplete_option(ctx: Context, args: t.List[str], param: Parameter) -> bool:
  377. """Determine if the given parameter is an option that needs a value.
  378. :param args: List of complete args before the incomplete value.
  379. :param param: Option object being checked.
  380. """
  381. if not isinstance(param, Option):
  382. return False
  383. if param.is_flag or param.count:
  384. return False
  385. last_option = None
  386. for index, arg in enumerate(reversed(args)):
  387. if index + 1 > param.nargs:
  388. break
  389. if _start_of_option(ctx, arg):
  390. last_option = arg
  391. return last_option is not None and last_option in param.opts
  392. def _resolve_context(
  393. cli: BaseCommand,
  394. ctx_args: t.MutableMapping[str, t.Any],
  395. prog_name: str,
  396. args: t.List[str],
  397. ) -> Context:
  398. """Produce the context hierarchy starting with the command and
  399. traversing the complete arguments. This only follows the commands,
  400. it doesn't trigger input prompts or callbacks.
  401. :param cli: Command being called.
  402. :param prog_name: Name of the executable in the shell.
  403. :param args: List of complete args before the incomplete value.
  404. """
  405. ctx_args["resilient_parsing"] = True
  406. ctx = cli.make_context(prog_name, args.copy(), **ctx_args)
  407. args = ctx.protected_args + ctx.args
  408. while args:
  409. command = ctx.command
  410. if isinstance(command, MultiCommand):
  411. if not command.chain:
  412. name, cmd, args = command.resolve_command(ctx, args)
  413. if cmd is None:
  414. return ctx
  415. ctx = cmd.make_context(name, args, parent=ctx, resilient_parsing=True)
  416. args = ctx.protected_args + ctx.args
  417. else:
  418. sub_ctx = ctx
  419. while args:
  420. name, cmd, args = command.resolve_command(ctx, args)
  421. if cmd is None:
  422. return ctx
  423. sub_ctx = cmd.make_context(
  424. name,
  425. args,
  426. parent=ctx,
  427. allow_extra_args=True,
  428. allow_interspersed_args=False,
  429. resilient_parsing=True,
  430. )
  431. args = sub_ctx.args
  432. ctx = sub_ctx
  433. args = [*sub_ctx.protected_args, *sub_ctx.args]
  434. else:
  435. break
  436. return ctx
  437. def _resolve_incomplete(
  438. ctx: Context, args: t.List[str], incomplete: str
  439. ) -> t.Tuple[t.Union[BaseCommand, Parameter], str]:
  440. """Find the Click object that will handle the completion of the
  441. incomplete value. Return the object and the incomplete value.
  442. :param ctx: Invocation context for the command represented by
  443. the parsed complete args.
  444. :param args: List of complete args before the incomplete value.
  445. :param incomplete: Value being completed. May be empty.
  446. """
  447. # Different shells treat an "=" between a long option name and
  448. # value differently. Might keep the value joined, return the "="
  449. # as a separate item, or return the split name and value. Always
  450. # split and discard the "=" to make completion easier.
  451. if incomplete == "=":
  452. incomplete = ""
  453. elif "=" in incomplete and _start_of_option(ctx, incomplete):
  454. name, _, incomplete = incomplete.partition("=")
  455. args.append(name)
  456. # The "--" marker tells Click to stop treating values as options
  457. # even if they start with the option character. If it hasn't been
  458. # given and the incomplete arg looks like an option, the current
  459. # command will provide option name completions.
  460. if "--" not in args and _start_of_option(ctx, incomplete):
  461. return ctx.command, incomplete
  462. params = ctx.command.get_params(ctx)
  463. # If the last complete arg is an option name with an incomplete
  464. # value, the option will provide value completions.
  465. for param in params:
  466. if _is_incomplete_option(ctx, args, param):
  467. return param, incomplete
  468. # It's not an option name or value. The first argument without a
  469. # parsed value will provide value completions.
  470. for param in params:
  471. if _is_incomplete_argument(ctx, param):
  472. return param, incomplete
  473. # There were no unparsed arguments, the command may be a group that
  474. # will provide command name completions.
  475. return ctx.command, incomplete