fuzzy_completer.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. from __future__ import annotations
  2. import re
  3. from typing import Callable, Iterable, NamedTuple
  4. from prompt_toolkit.document import Document
  5. from prompt_toolkit.filters import FilterOrBool, to_filter
  6. from prompt_toolkit.formatted_text import AnyFormattedText, StyleAndTextTuples
  7. from .base import CompleteEvent, Completer, Completion
  8. from .word_completer import WordCompleter
  9. __all__ = [
  10. "FuzzyCompleter",
  11. "FuzzyWordCompleter",
  12. ]
  13. class FuzzyCompleter(Completer):
  14. """
  15. Fuzzy completion.
  16. This wraps any other completer and turns it into a fuzzy completer.
  17. If the list of words is: ["leopard" , "gorilla", "dinosaur", "cat", "bee"]
  18. Then trying to complete "oar" would yield "leopard" and "dinosaur", but not
  19. the others, because they match the regular expression 'o.*a.*r'.
  20. Similar, in another application "djm" could expand to "django_migrations".
  21. The results are sorted by relevance, which is defined as the start position
  22. and the length of the match.
  23. Notice that this is not really a tool to work around spelling mistakes,
  24. like what would be possible with difflib. The purpose is rather to have a
  25. quicker or more intuitive way to filter the given completions, especially
  26. when many completions have a common prefix.
  27. Fuzzy algorithm is based on this post:
  28. https://blog.amjith.com/fuzzyfinder-in-10-lines-of-python
  29. :param completer: A :class:`~.Completer` instance.
  30. :param WORD: When True, use WORD characters.
  31. :param pattern: Regex pattern which selects the characters before the
  32. cursor that are considered for the fuzzy matching.
  33. :param enable_fuzzy: (bool or `Filter`) Enabled the fuzzy behavior. For
  34. easily turning fuzzyness on or off according to a certain condition.
  35. """
  36. def __init__(
  37. self,
  38. completer: Completer,
  39. WORD: bool = False,
  40. pattern: str | None = None,
  41. enable_fuzzy: FilterOrBool = True,
  42. ) -> None:
  43. assert pattern is None or pattern.startswith("^")
  44. self.completer = completer
  45. self.pattern = pattern
  46. self.WORD = WORD
  47. self.pattern = pattern
  48. self.enable_fuzzy = to_filter(enable_fuzzy)
  49. def get_completions(
  50. self, document: Document, complete_event: CompleteEvent
  51. ) -> Iterable[Completion]:
  52. if self.enable_fuzzy():
  53. return self._get_fuzzy_completions(document, complete_event)
  54. else:
  55. return self.completer.get_completions(document, complete_event)
  56. def _get_pattern(self) -> str:
  57. if self.pattern:
  58. return self.pattern
  59. if self.WORD:
  60. return r"[^\s]+"
  61. return "^[a-zA-Z0-9_]*"
  62. def _get_fuzzy_completions(
  63. self, document: Document, complete_event: CompleteEvent
  64. ) -> Iterable[Completion]:
  65. word_before_cursor = document.get_word_before_cursor(
  66. pattern=re.compile(self._get_pattern())
  67. )
  68. # Get completions
  69. document2 = Document(
  70. text=document.text[: document.cursor_position - len(word_before_cursor)],
  71. cursor_position=document.cursor_position - len(word_before_cursor),
  72. )
  73. inner_completions = list(
  74. self.completer.get_completions(document2, complete_event)
  75. )
  76. fuzzy_matches: list[_FuzzyMatch] = []
  77. if word_before_cursor == "":
  78. # If word before the cursor is an empty string, consider all
  79. # completions, without filtering everything with an empty regex
  80. # pattern.
  81. fuzzy_matches = [_FuzzyMatch(0, 0, compl) for compl in inner_completions]
  82. else:
  83. pat = ".*?".join(map(re.escape, word_before_cursor))
  84. pat = f"(?=({pat}))" # lookahead regex to manage overlapping matches
  85. regex = re.compile(pat, re.IGNORECASE)
  86. for compl in inner_completions:
  87. matches = list(regex.finditer(compl.text))
  88. if matches:
  89. # Prefer the match, closest to the left, then shortest.
  90. best = min(matches, key=lambda m: (m.start(), len(m.group(1))))
  91. fuzzy_matches.append(
  92. _FuzzyMatch(len(best.group(1)), best.start(), compl)
  93. )
  94. def sort_key(fuzzy_match: _FuzzyMatch) -> tuple[int, int]:
  95. "Sort by start position, then by the length of the match."
  96. return fuzzy_match.start_pos, fuzzy_match.match_length
  97. fuzzy_matches = sorted(fuzzy_matches, key=sort_key)
  98. for match in fuzzy_matches:
  99. # Include these completions, but set the correct `display`
  100. # attribute and `start_position`.
  101. yield Completion(
  102. text=match.completion.text,
  103. start_position=match.completion.start_position
  104. - len(word_before_cursor),
  105. # We access to private `_display_meta` attribute, because that one is lazy.
  106. display_meta=match.completion._display_meta,
  107. display=self._get_display(match, word_before_cursor),
  108. style=match.completion.style,
  109. )
  110. def _get_display(
  111. self, fuzzy_match: _FuzzyMatch, word_before_cursor: str
  112. ) -> AnyFormattedText:
  113. """
  114. Generate formatted text for the display label.
  115. """
  116. def get_display() -> AnyFormattedText:
  117. m = fuzzy_match
  118. word = m.completion.text
  119. if m.match_length == 0:
  120. # No highlighting when we have zero length matches (no input text).
  121. # In this case, use the original display text (which can include
  122. # additional styling or characters).
  123. return m.completion.display
  124. result: StyleAndTextTuples = []
  125. # Text before match.
  126. result.append(("class:fuzzymatch.outside", word[: m.start_pos]))
  127. # The match itself.
  128. characters = list(word_before_cursor)
  129. for c in word[m.start_pos : m.start_pos + m.match_length]:
  130. classname = "class:fuzzymatch.inside"
  131. if characters and c.lower() == characters[0].lower():
  132. classname += ".character"
  133. del characters[0]
  134. result.append((classname, c))
  135. # Text after match.
  136. result.append(
  137. ("class:fuzzymatch.outside", word[m.start_pos + m.match_length :])
  138. )
  139. return result
  140. return get_display()
  141. class FuzzyWordCompleter(Completer):
  142. """
  143. Fuzzy completion on a list of words.
  144. (This is basically a `WordCompleter` wrapped in a `FuzzyCompleter`.)
  145. :param words: List of words or callable that returns a list of words.
  146. :param meta_dict: Optional dict mapping words to their meta-information.
  147. :param WORD: When True, use WORD characters.
  148. """
  149. def __init__(
  150. self,
  151. words: list[str] | Callable[[], list[str]],
  152. meta_dict: dict[str, str] | None = None,
  153. WORD: bool = False,
  154. ) -> None:
  155. self.words = words
  156. self.meta_dict = meta_dict or {}
  157. self.WORD = WORD
  158. self.word_completer = WordCompleter(
  159. words=self.words, WORD=self.WORD, meta_dict=self.meta_dict
  160. )
  161. self.fuzzy_completer = FuzzyCompleter(self.word_completer, WORD=self.WORD)
  162. def get_completions(
  163. self, document: Document, complete_event: CompleteEvent
  164. ) -> Iterable[Completion]:
  165. return self.fuzzy_completer.get_completions(document, complete_event)
  166. class _FuzzyMatch(NamedTuple):
  167. match_length: int
  168. start_pos: int
  169. completion: Completion