word_completer.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. from __future__ import annotations
  2. from typing import Callable, Iterable, Mapping, Pattern
  3. from prompt_toolkit.completion import CompleteEvent, Completer, Completion
  4. from prompt_toolkit.document import Document
  5. from prompt_toolkit.formatted_text import AnyFormattedText
  6. __all__ = [
  7. "WordCompleter",
  8. ]
  9. class WordCompleter(Completer):
  10. """
  11. Simple autocompletion on a list of words.
  12. :param words: List of words or callable that returns a list of words.
  13. :param ignore_case: If True, case-insensitive completion.
  14. :param meta_dict: Optional dict mapping words to their meta-text. (This
  15. should map strings to strings or formatted text.)
  16. :param WORD: When True, use WORD characters.
  17. :param sentence: When True, don't complete by comparing the word before the
  18. cursor, but by comparing all the text before the cursor. In this case,
  19. the list of words is just a list of strings, where each string can
  20. contain spaces. (Can not be used together with the WORD option.)
  21. :param match_middle: When True, match not only the start, but also in the
  22. middle of the word.
  23. :param pattern: Optional compiled regex for finding the word before
  24. the cursor to complete. When given, use this regex pattern instead of
  25. default one (see document._FIND_WORD_RE)
  26. """
  27. def __init__(
  28. self,
  29. words: list[str] | Callable[[], list[str]],
  30. ignore_case: bool = False,
  31. display_dict: Mapping[str, AnyFormattedText] | None = None,
  32. meta_dict: Mapping[str, AnyFormattedText] | None = None,
  33. WORD: bool = False,
  34. sentence: bool = False,
  35. match_middle: bool = False,
  36. pattern: Pattern[str] | None = None,
  37. ) -> None:
  38. assert not (WORD and sentence)
  39. self.words = words
  40. self.ignore_case = ignore_case
  41. self.display_dict = display_dict or {}
  42. self.meta_dict = meta_dict or {}
  43. self.WORD = WORD
  44. self.sentence = sentence
  45. self.match_middle = match_middle
  46. self.pattern = pattern
  47. def get_completions(
  48. self, document: Document, complete_event: CompleteEvent
  49. ) -> Iterable[Completion]:
  50. # Get list of words.
  51. words = self.words
  52. if callable(words):
  53. words = words()
  54. # Get word/text before cursor.
  55. if self.sentence:
  56. word_before_cursor = document.text_before_cursor
  57. else:
  58. word_before_cursor = document.get_word_before_cursor(
  59. WORD=self.WORD, pattern=self.pattern
  60. )
  61. if self.ignore_case:
  62. word_before_cursor = word_before_cursor.lower()
  63. def word_matches(word: str) -> bool:
  64. """True when the word before the cursor matches."""
  65. if self.ignore_case:
  66. word = word.lower()
  67. if self.match_middle:
  68. return word_before_cursor in word
  69. else:
  70. return word.startswith(word_before_cursor)
  71. for a in words:
  72. if word_matches(a):
  73. display = self.display_dict.get(a, a)
  74. display_meta = self.meta_dict.get(a, "")
  75. yield Completion(
  76. text=a,
  77. start_position=-len(word_before_cursor),
  78. display=display,
  79. display_meta=display_meta,
  80. )