1# Copyright 2012-2023, Andrey Kislyuk and argcomplete contributors. Licensed under the terms of the
2# `Apache License, Version 2.0 <http://www.apache.org/licenses/LICENSE-2.0>`_. Distribution of the LICENSE and NOTICE
3# files with source copies of this package and derivative works is **REQUIRED** as specified by the Apache License.
4# See https://github.com/kislyuk/argcomplete for more info.
5
6import argparse
7import os
8import sys
9from collections.abc import Mapping
10from typing import Callable, Dict, List, Optional, Sequence, TextIO, Union
11
12from . import io as _io
13from .completers import BaseCompleter, ChoicesCompleter, FilesCompleter, SuppressCompleter
14from .io import debug, mute_stderr
15from .lexers import split_line
16from .packages._argparse import IntrospectiveArgumentParser, action_is_greedy, action_is_open, action_is_satisfied
17
18safe_actions = {
19 argparse._StoreAction,
20 argparse._StoreConstAction,
21 argparse._StoreTrueAction,
22 argparse._StoreFalseAction,
23 argparse._AppendAction,
24 argparse._AppendConstAction,
25 argparse._CountAction,
26}
27
28
29def default_validator(completion, prefix):
30 return completion.startswith(prefix)
31
32
33class CompletionFinder(object):
34 """
35 Inherit from this class if you wish to override any of the stages below. Otherwise, use
36 ``argcomplete.autocomplete()`` directly (it's a convenience instance of this class). It has the same signature as
37 :meth:`CompletionFinder.__call__()`.
38 """
39
40 def __init__(
41 self,
42 argument_parser=None,
43 always_complete_options=True,
44 exclude=None,
45 validator=None,
46 print_suppressed=False,
47 default_completer=FilesCompleter(),
48 append_space=None,
49 ):
50 self._parser = argument_parser
51 self._formatter = None
52 self.always_complete_options = always_complete_options
53 self.exclude = exclude
54 if validator is None:
55 validator = default_validator
56 self.validator = validator
57 self.print_suppressed = print_suppressed
58 self.completing = False
59 self._display_completions: Dict[str, str] = {}
60 self.default_completer = default_completer
61 if append_space is None:
62 append_space = os.environ.get("_ARGCOMPLETE_SUPPRESS_SPACE") != "1"
63 self.append_space = append_space
64
65 def __call__(
66 self,
67 argument_parser: argparse.ArgumentParser,
68 always_complete_options: Union[bool, str] = True,
69 exit_method: Callable = os._exit,
70 output_stream: Optional[TextIO] = None,
71 exclude: Optional[Sequence[str]] = None,
72 validator: Optional[Callable] = None,
73 print_suppressed: bool = False,
74 append_space: Optional[bool] = None,
75 default_completer: BaseCompleter = FilesCompleter(),
76 ) -> None:
77 """
78 :param argument_parser: The argument parser to autocomplete on
79 :param always_complete_options:
80 Controls the autocompletion of option strings if an option string opening character (normally ``-``) has not
81 been entered. If ``True`` (default), both short (``-x``) and long (``--x``) option strings will be
82 suggested. If ``False``, no option strings will be suggested. If ``long``, long options and short options
83 with no long variant will be suggested. If ``short``, short options and long options with no short variant
84 will be suggested.
85 :param exit_method:
86 Method used to stop the program after printing completions. Defaults to :meth:`os._exit`. If you want to
87 perform a normal exit that calls exit handlers, use :meth:`sys.exit`.
88 :param exclude: List of strings representing options to be omitted from autocompletion
89 :param validator:
90 Function to filter all completions through before returning (called with two string arguments, completion
91 and prefix; return value is evaluated as a boolean)
92 :param print_suppressed:
93 Whether or not to autocomplete options that have the ``help=argparse.SUPPRESS`` keyword argument set.
94 :param append_space:
95 Whether to append a space to unique matches. The default is ``True``.
96
97 .. note::
98 If you are not subclassing CompletionFinder to override its behaviors,
99 use :meth:`argcomplete.autocomplete()` directly. It has the same signature as this method.
100
101 Produces tab completions for ``argument_parser``. See module docs for more info.
102
103 Argcomplete only executes actions if their class is known not to have side effects. Custom action classes can be
104 added to argcomplete.safe_actions, if their values are wanted in the ``parsed_args`` completer argument, or
105 their execution is otherwise desirable.
106 """
107 self.__init__( # type: ignore
108 argument_parser,
109 always_complete_options=always_complete_options,
110 exclude=exclude,
111 validator=validator,
112 print_suppressed=print_suppressed,
113 append_space=append_space,
114 default_completer=default_completer,
115 )
116
117 if "_ARGCOMPLETE" not in os.environ:
118 # not an argument completion invocation
119 return
120
121 self._init_debug_stream()
122
123 if output_stream is None:
124 filename = os.environ.get("_ARGCOMPLETE_STDOUT_FILENAME")
125 if filename is not None:
126 debug("Using output file {}".format(filename))
127 output_stream = open(filename, "w")
128
129 if output_stream is None:
130 try:
131 output_stream = os.fdopen(8, "w")
132 except Exception:
133 debug("Unable to open fd 8 for writing, quitting")
134 exit_method(1)
135
136 assert output_stream is not None
137
138 ifs = os.environ.get("_ARGCOMPLETE_IFS", "\013")
139 if len(ifs) != 1:
140 debug("Invalid value for IFS, quitting [{v}]".format(v=ifs))
141 exit_method(1)
142
143 dfs = os.environ.get("_ARGCOMPLETE_DFS")
144 if dfs and len(dfs) != 1:
145 debug("Invalid value for DFS, quitting [{v}]".format(v=dfs))
146 exit_method(1)
147
148 comp_line = os.environ["COMP_LINE"]
149 comp_point = int(os.environ["COMP_POINT"])
150
151 cword_prequote, cword_prefix, cword_suffix, comp_words, last_wordbreak_pos = split_line(comp_line, comp_point)
152
153 # _ARGCOMPLETE is set by the shell script to tell us where comp_words
154 # should start, based on what we're completing.
155 # 1: <script> [args]
156 # 2: python <script> [args]
157 # 3: python -m <module> [args]
158 start = int(os.environ["_ARGCOMPLETE"]) - 1
159 comp_words = comp_words[start:]
160
161 if cword_prefix and cword_prefix[0] in self._parser.prefix_chars and "=" in cword_prefix:
162 # Special case for when the current word is "--optional=PARTIAL_VALUE". Give the optional to the parser.
163 comp_words.append(cword_prefix.split("=", 1)[0])
164
165 debug(
166 "\nLINE: {!r}".format(comp_line),
167 "\nPOINT: {!r}".format(comp_point),
168 "\nPREQUOTE: {!r}".format(cword_prequote),
169 "\nPREFIX: {!r}".format(cword_prefix),
170 "\nSUFFIX: {!r}".format(cword_suffix),
171 "\nWORDS:",
172 comp_words,
173 )
174
175 completions = self._get_completions(comp_words, cword_prefix, cword_prequote, last_wordbreak_pos)
176
177 if dfs:
178 display_completions = {
179 key: value.replace(ifs, " ") if value else "" for key, value in self._display_completions.items()
180 }
181 completions = [dfs.join((key, display_completions.get(key) or "")) for key in completions]
182
183 if os.environ.get("_ARGCOMPLETE_SHELL") == "zsh":
184 completions = [f"{c}:{self._display_completions.get(c)}" for c in completions]
185
186 debug("\nReturning completions:", completions)
187 output_stream.write(ifs.join(completions))
188 output_stream.flush()
189 _io.debug_stream.flush()
190 exit_method(0)
191
192 def _init_debug_stream(self):
193 """Initialize debug output stream
194
195 By default, writes to file descriptor 9, or stderr if that fails.
196 This can be overridden by derived classes, for example to avoid
197 clashes with file descriptors being used elsewhere (such as in pytest).
198 """
199 try:
200 _io.debug_stream = os.fdopen(9, "w")
201 except Exception:
202 _io.debug_stream = sys.stderr
203 debug()
204
205 def _get_completions(self, comp_words, cword_prefix, cword_prequote, last_wordbreak_pos):
206 active_parsers = self._patch_argument_parser()
207
208 parsed_args = argparse.Namespace()
209 self.completing = True
210
211 try:
212 debug("invoking parser with", comp_words[1:])
213 with mute_stderr():
214 a = self._parser.parse_known_args(comp_words[1:], namespace=parsed_args)
215 debug("parsed args:", a)
216 except BaseException as e:
217 debug("\nexception", type(e), str(e), "while parsing args")
218
219 self.completing = False
220
221 if "--" in comp_words:
222 self.always_complete_options = False
223
224 completions = self.collect_completions(active_parsers, parsed_args, cword_prefix)
225 completions = self.filter_completions(completions)
226 completions = self.quote_completions(completions, cword_prequote, last_wordbreak_pos)
227 return completions
228
229 def _patch_argument_parser(self):
230 """
231 Since argparse doesn't support much introspection, we monkey-patch it to replace the parse_known_args method and
232 all actions with hooks that tell us which action was last taken or about to be taken, and let us have the parser
233 figure out which subparsers need to be activated (then recursively monkey-patch those).
234 We save all active ArgumentParsers to extract all their possible option names later.
235 """
236 self.active_parsers: List[argparse.ArgumentParser] = []
237 self.visited_positionals: List[argparse.Action] = []
238
239 completer = self
240
241 def patch(parser):
242 completer.visited_positionals.append(parser)
243 completer.active_parsers.append(parser)
244
245 if isinstance(parser, IntrospectiveArgumentParser):
246 return
247
248 classname = "MonkeyPatchedIntrospectiveArgumentParser"
249
250 parser.__class__ = type(classname, (IntrospectiveArgumentParser, parser.__class__), {})
251
252 for action in parser._actions:
253 if hasattr(action, "_orig_class"):
254 continue
255
256 # TODO: accomplish this with super
257 class IntrospectAction(action.__class__): # type: ignore
258 def __call__(self, parser, namespace, values, option_string=None):
259 debug("Action stub called on", self)
260 debug("\targs:", parser, namespace, values, option_string)
261 debug("\torig class:", self._orig_class)
262 debug("\torig callable:", self._orig_callable)
263
264 if not completer.completing:
265 self._orig_callable(parser, namespace, values, option_string=option_string)
266 elif issubclass(self._orig_class, argparse._SubParsersAction):
267 debug("orig class is a subparsers action: patching and running it")
268 patch(self._name_parser_map[values[0]])
269 self._orig_callable(parser, namespace, values, option_string=option_string)
270 elif self._orig_class in safe_actions:
271 if not self.option_strings:
272 completer.visited_positionals.append(self)
273
274 self._orig_callable(parser, namespace, values, option_string=option_string)
275
276 action._orig_class = action.__class__
277 action._orig_callable = action.__call__
278 action.__class__ = IntrospectAction
279
280 patch(self._parser)
281
282 debug("Active parsers:", self.active_parsers)
283 debug("Visited positionals:", self.visited_positionals)
284
285 return self.active_parsers
286
287 def _get_action_help(self, action):
288 if action.help is None:
289 return ""
290 if "%" not in action.help:
291 return action.help
292 if self._formatter is None:
293 self._formatter = self._parser.formatter_class(prog=self._parser.prog)
294 return self._formatter._expand_help(action)
295
296 def _get_subparser_completions(self, parser, cword_prefix):
297 aliases_by_parser: Dict[argparse.ArgumentParser, List[str]] = {}
298 for key in parser.choices.keys():
299 p = parser.choices[key]
300 aliases_by_parser.setdefault(p, []).append(key)
301
302 for action in parser._get_subactions():
303 for alias in aliases_by_parser[parser.choices[action.dest]]:
304 if alias.startswith(cword_prefix):
305 self._display_completions[alias] = self._get_action_help(action)
306
307 completions = [subcmd for subcmd in parser.choices.keys() if subcmd.startswith(cword_prefix)]
308 return completions
309
310 def _include_options(self, action, cword_prefix):
311 if len(cword_prefix) > 0 or self.always_complete_options is True:
312 return [opt for opt in action.option_strings if opt.startswith(cword_prefix)]
313 long_opts = [opt for opt in action.option_strings if len(opt) > 2]
314 short_opts = [opt for opt in action.option_strings if len(opt) <= 2]
315 if self.always_complete_options == "long":
316 return long_opts if long_opts else short_opts
317 elif self.always_complete_options == "short":
318 return short_opts if short_opts else long_opts
319 return []
320
321 def _get_option_completions(self, parser, cword_prefix):
322 for action in parser._actions:
323 if action.option_strings:
324 for option_string in action.option_strings:
325 if option_string.startswith(cword_prefix):
326 self._display_completions[option_string] = self._get_action_help(action)
327
328 option_completions = []
329 for action in parser._actions:
330 if not self.print_suppressed:
331 completer = getattr(action, "completer", None)
332 if isinstance(completer, SuppressCompleter) and completer.suppress():
333 continue
334 if action.help == argparse.SUPPRESS:
335 continue
336 if not self._action_allowed(action, parser):
337 continue
338 if not isinstance(action, argparse._SubParsersAction):
339 option_completions += self._include_options(action, cword_prefix)
340 return option_completions
341
342 @staticmethod
343 def _action_allowed(action, parser):
344 # Logic adapted from take_action in ArgumentParser._parse_known_args
345 # (members are saved by vendor._argparse.IntrospectiveArgumentParser)
346 for conflict_action in parser._action_conflicts.get(action, []):
347 if conflict_action in parser._seen_non_default_actions:
348 return False
349 return True
350
351 def _complete_active_option(self, parser, next_positional, cword_prefix, parsed_args, completions):
352 debug("Active actions (L={l}): {a}".format(l=len(parser.active_actions), a=parser.active_actions))
353
354 isoptional = cword_prefix and cword_prefix[0] in parser.prefix_chars
355 optional_prefix = ""
356 greedy_actions = [x for x in parser.active_actions if action_is_greedy(x, isoptional)]
357 if greedy_actions:
358 assert len(greedy_actions) == 1, "expect at most 1 greedy action"
359 # This means the action will fail to parse if the word under the cursor is not given
360 # to it, so give it exclusive control over completions (flush previous completions)
361 debug("Resetting completions because", greedy_actions[0], "must consume the next argument")
362 self._display_completions = {}
363 completions = []
364 elif isoptional:
365 if "=" in cword_prefix:
366 # Special case for when the current word is "--optional=PARTIAL_VALUE".
367 # The completer runs on PARTIAL_VALUE. The prefix is added back to the completions
368 # (and chopped back off later in quote_completions() by the COMP_WORDBREAKS logic).
369 optional_prefix, _, cword_prefix = cword_prefix.partition("=")
370 else:
371 # Only run completers if current word does not start with - (is not an optional)
372 return completions
373
374 complete_remaining_positionals = False
375 # Use the single greedy action (if there is one) or all active actions.
376 for active_action in greedy_actions or parser.active_actions:
377 if not active_action.option_strings: # action is a positional
378 if action_is_open(active_action):
379 # Any positional arguments after this may slide down into this action
380 # if more arguments are added (since the user may not be done yet),
381 # so it is extremely difficult to tell which completers to run.
382 # Running all remaining completers will probably show more than the user wants
383 # but it also guarantees we won't miss anything.
384 complete_remaining_positionals = True
385 if not complete_remaining_positionals:
386 if action_is_satisfied(active_action) and not action_is_open(active_action):
387 debug("Skipping", active_action)
388 continue
389
390 debug("Activating completion for", active_action, active_action._orig_class)
391 # completer = getattr(active_action, "completer", DefaultCompleter())
392 completer = getattr(active_action, "completer", None)
393
394 if completer is None:
395 if active_action.choices is not None and not isinstance(active_action, argparse._SubParsersAction):
396 completer = ChoicesCompleter(active_action.choices)
397 elif not isinstance(active_action, argparse._SubParsersAction):
398 completer = self.default_completer
399
400 if completer:
401 if isinstance(completer, SuppressCompleter) and completer.suppress():
402 continue
403
404 if callable(completer):
405 completer_output = completer(
406 prefix=cword_prefix, action=active_action, parser=parser, parsed_args=parsed_args
407 )
408 if isinstance(completer_output, Mapping):
409 for completion, description in completer_output.items():
410 if self.validator(completion, cword_prefix):
411 completions.append(completion)
412 self._display_completions[completion] = description
413 else:
414 for completion in completer_output:
415 if self.validator(completion, cword_prefix):
416 completions.append(completion)
417 if isinstance(completer, ChoicesCompleter):
418 self._display_completions[completion] = self._get_action_help(action)
419 else:
420 self._display_completions[completion] = ""
421 else:
422 debug("Completer is not callable, trying the readline completer protocol instead")
423 for i in range(9999):
424 next_completion = completer.complete(cword_prefix, i) # type: ignore
425 if next_completion is None:
426 break
427 if self.validator(next_completion, cword_prefix):
428 self._display_completions[next_completion] = ""
429 completions.append(next_completion)
430 if optional_prefix:
431 completions = [optional_prefix + "=" + completion for completion in completions]
432 debug("Completions:", completions)
433 return completions
434
435 def collect_completions(
436 self, active_parsers: List[argparse.ArgumentParser], parsed_args: argparse.Namespace, cword_prefix: str
437 ) -> List[str]:
438 """
439 Visits the active parsers and their actions, executes their completers or introspects them to collect their
440 option strings. Returns the resulting completions as a list of strings.
441
442 This method is exposed for overriding in subclasses; there is no need to use it directly.
443 """
444 completions: List[str] = []
445
446 debug("all active parsers:", active_parsers)
447 active_parser = active_parsers[-1]
448 debug("active_parser:", active_parser)
449 if self.always_complete_options or (len(cword_prefix) > 0 and cword_prefix[0] in active_parser.prefix_chars):
450 completions += self._get_option_completions(active_parser, cword_prefix)
451 debug("optional options:", completions)
452
453 next_positional = self._get_next_positional()
454 debug("next_positional:", next_positional)
455
456 if isinstance(next_positional, argparse._SubParsersAction):
457 completions += self._get_subparser_completions(next_positional, cword_prefix)
458
459 completions = self._complete_active_option(
460 active_parser, next_positional, cword_prefix, parsed_args, completions
461 )
462 debug("active options:", completions)
463 debug("display completions:", self._display_completions)
464
465 return completions
466
467 def _get_next_positional(self):
468 """
469 Get the next positional action if it exists.
470 """
471 active_parser = self.active_parsers[-1]
472 last_positional = self.visited_positionals[-1]
473
474 all_positionals = active_parser._get_positional_actions()
475 if not all_positionals:
476 return None
477
478 if active_parser == last_positional:
479 return all_positionals[0]
480
481 i = 0
482 for i in range(len(all_positionals)):
483 if all_positionals[i] == last_positional:
484 break
485
486 if i + 1 < len(all_positionals):
487 return all_positionals[i + 1]
488
489 return None
490
491 def filter_completions(self, completions: List[str]) -> List[str]:
492 """
493 De-duplicates completions and excludes those specified by ``exclude``.
494 Returns the filtered completions as a list.
495
496 This method is exposed for overriding in subclasses; there is no need to use it directly.
497 """
498 filtered_completions = []
499 for completion in completions:
500 if self.exclude is not None:
501 if completion in self.exclude:
502 continue
503 if completion not in filtered_completions:
504 filtered_completions.append(completion)
505 return filtered_completions
506
507 def quote_completions(
508 self, completions: List[str], cword_prequote: str, last_wordbreak_pos: Optional[int]
509 ) -> List[str]:
510 """
511 If the word under the cursor started with a quote (as indicated by a nonempty ``cword_prequote``), escapes
512 occurrences of that quote character in the completions, and adds the quote to the beginning of each completion.
513 Otherwise, escapes all characters that bash splits words on (``COMP_WORDBREAKS``), and removes portions of
514 completions before the first colon if (``COMP_WORDBREAKS``) contains a colon.
515
516 If there is only one completion, and it doesn't end with a **continuation character** (``/``, ``:``, or ``=``),
517 adds a space after the completion.
518
519 This method is exposed for overriding in subclasses; there is no need to use it directly.
520 """
521 special_chars = "\\"
522 # If the word under the cursor was quoted, escape the quote char.
523 # Otherwise, escape all special characters and specially handle all COMP_WORDBREAKS chars.
524 if cword_prequote == "":
525 # Bash mangles completions which contain characters in COMP_WORDBREAKS.
526 # This workaround has the same effect as __ltrim_colon_completions in bash_completion
527 # (extended to characters other than the colon).
528 if last_wordbreak_pos is not None:
529 completions = [c[last_wordbreak_pos + 1 :] for c in completions]
530 special_chars += "();<>|&!`$* \t\n\"'"
531 elif cword_prequote == '"':
532 special_chars += '"`$!'
533
534 if os.environ.get("_ARGCOMPLETE_SHELL") in ("tcsh", "fish"):
535 # tcsh and fish escapes special characters itself.
536 special_chars = ""
537 elif cword_prequote == "'":
538 # Nothing can be escaped in single quotes, so we need to close
539 # the string, escape the single quote, then open a new string.
540 special_chars = ""
541 completions = [c.replace("'", r"'\''") for c in completions]
542
543 # PowerShell uses ` as escape character.
544 if os.environ.get("_ARGCOMPLETE_SHELL") == "powershell":
545 escape_char = '`'
546 special_chars = special_chars.replace('`', '')
547 else:
548 escape_char = "\\"
549 if os.environ.get("_ARGCOMPLETE_SHELL") == "zsh":
550 # zsh uses colon as a separator between a completion and its description.
551 special_chars += ":"
552
553 escaped_completions = []
554 for completion in completions:
555 escaped_completion = completion
556 for char in special_chars:
557 escaped_completion = escaped_completion.replace(char, escape_char + char)
558 escaped_completions.append(escaped_completion)
559 if completion in self._display_completions:
560 self._display_completions[escaped_completion] = self._display_completions[completion]
561
562 if self.append_space:
563 # Similar functionality in bash was previously turned off by supplying the "-o nospace" option to complete.
564 # Now it is conditionally disabled using "compopt -o nospace" if the match ends in a continuation character.
565 # This code is retained for environments where this isn't done natively.
566 continuation_chars = "=/:"
567 if len(escaped_completions) == 1 and escaped_completions[0][-1] not in continuation_chars:
568 if cword_prequote == "":
569 escaped_completions[0] += " "
570
571 return escaped_completions
572
573 def rl_complete(self, text, state):
574 """
575 Alternate entry point for using the argcomplete completer in a readline-based REPL. See also
576 `rlcompleter <https://docs.python.org/3/library/rlcompleter.html#completer-objects>`_.
577 Usage:
578
579 .. code-block:: python
580
581 import argcomplete, argparse, readline
582 parser = argparse.ArgumentParser()
583 ...
584 completer = argcomplete.CompletionFinder(parser)
585 readline.set_completer_delims("")
586 readline.set_completer(completer.rl_complete)
587 readline.parse_and_bind("tab: complete")
588 result = input("prompt> ")
589 """
590 if state == 0:
591 cword_prequote, cword_prefix, cword_suffix, comp_words, first_colon_pos = split_line(text)
592 comp_words.insert(0, sys.argv[0])
593 matches = self._get_completions(comp_words, cword_prefix, cword_prequote, first_colon_pos)
594 self._rl_matches = [text + match[len(cword_prefix) :] for match in matches]
595
596 if state < len(self._rl_matches):
597 return self._rl_matches[state]
598 else:
599 return None
600
601 def get_display_completions(self):
602 """
603 This function returns a mapping of completions to their help strings for displaying to the user.
604 """
605 return self._display_completions
606
607
608class ExclusiveCompletionFinder(CompletionFinder):
609 @staticmethod
610 def _action_allowed(action, parser):
611 if not CompletionFinder._action_allowed(action, parser):
612 return False
613
614 append_classes = (argparse._AppendAction, argparse._AppendConstAction)
615 if action._orig_class in append_classes:
616 return True
617
618 if action not in parser._seen_non_default_actions:
619 return True
620
621 return False