completion.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. """
  2. Completer for a regular grammar.
  3. """
  4. from __future__ import annotations
  5. from typing import Iterable
  6. from prompt_toolkit.completion import CompleteEvent, Completer, Completion
  7. from prompt_toolkit.document import Document
  8. from .compiler import Match, _CompiledGrammar
  9. __all__ = [
  10. "GrammarCompleter",
  11. ]
  12. class GrammarCompleter(Completer):
  13. """
  14. Completer which can be used for autocompletion according to variables in
  15. the grammar. Each variable can have a different autocompleter.
  16. :param compiled_grammar: `GrammarCompleter` instance.
  17. :param completers: `dict` mapping variable names of the grammar to the
  18. `Completer` instances to be used for each variable.
  19. """
  20. def __init__(
  21. self, compiled_grammar: _CompiledGrammar, completers: dict[str, Completer]
  22. ) -> None:
  23. self.compiled_grammar = compiled_grammar
  24. self.completers = completers
  25. def get_completions(
  26. self, document: Document, complete_event: CompleteEvent
  27. ) -> Iterable[Completion]:
  28. m = self.compiled_grammar.match_prefix(document.text_before_cursor)
  29. if m:
  30. yield from self._remove_duplicates(
  31. self._get_completions_for_match(m, complete_event)
  32. )
  33. def _get_completions_for_match(
  34. self, match: Match, complete_event: CompleteEvent
  35. ) -> Iterable[Completion]:
  36. """
  37. Yield all the possible completions for this input string.
  38. (The completer assumes that the cursor position was at the end of the
  39. input string.)
  40. """
  41. for match_variable in match.end_nodes():
  42. varname = match_variable.varname
  43. start = match_variable.start
  44. completer = self.completers.get(varname)
  45. if completer:
  46. text = match_variable.value
  47. # Unwrap text.
  48. unwrapped_text = self.compiled_grammar.unescape(varname, text)
  49. # Create a document, for the completions API (text/cursor_position)
  50. document = Document(unwrapped_text, len(unwrapped_text))
  51. # Call completer
  52. for completion in completer.get_completions(document, complete_event):
  53. new_text = (
  54. unwrapped_text[: len(text) + completion.start_position]
  55. + completion.text
  56. )
  57. # Wrap again.
  58. yield Completion(
  59. text=self.compiled_grammar.escape(varname, new_text),
  60. start_position=start - len(match.string),
  61. display=completion.display,
  62. display_meta=completion.display_meta,
  63. )
  64. def _remove_duplicates(self, items: Iterable[Completion]) -> Iterable[Completion]:
  65. """
  66. Remove duplicates, while keeping the order.
  67. (Sometimes we have duplicates, because the there several matches of the
  68. same grammar, each yielding similar completions.)
  69. """
  70. def hash_completion(completion: Completion) -> tuple[str, int]:
  71. return completion.text, completion.start_position
  72. yielded_so_far: set[tuple[str, int]] = set()
  73. for completion in items:
  74. hash_value = hash_completion(completion)
  75. if hash_value not in yielded_so_far:
  76. yielded_so_far.add(hash_value)
  77. yield completion