From 7775081fa2e64f468f414684757d0ab1ee0d7031 Mon Sep 17 00:00:00 2001 From: fan wu <336798967+FanWu-ai@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:45:53 +0900 Subject: [PATCH] Preserve asynchronous completion in FuzzyCompleter --- .../completion/fuzzy_completer.py | 41 ++++- tests/test_fuzzy_async.py | 151 ++++++++++++++++++ 2 files changed, 185 insertions(+), 7 deletions(-) create mode 100644 tests/test_fuzzy_async.py diff --git a/src/prompt_toolkit/completion/fuzzy_completer.py b/src/prompt_toolkit/completion/fuzzy_completer.py index 0f8bba740..9ffb8be38 100644 --- a/src/prompt_toolkit/completion/fuzzy_completer.py +++ b/src/prompt_toolkit/completion/fuzzy_completer.py @@ -1,7 +1,8 @@ from __future__ import annotations import re -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence +from contextlib import aclosing from typing import NamedTuple from prompt_toolkit.document import Document @@ -76,23 +77,49 @@ def _get_pattern(self) -> str: return r"[^\s]+" return "^[a-zA-Z0-9_]*" - def _get_fuzzy_completions( + async def get_completions_async( self, document: Document, complete_event: CompleteEvent - ) -> Iterable[Completion]: + ) -> AsyncGenerator[Completion, None]: + if self.enable_fuzzy(): + word_before_cursor, document2 = self._get_fuzzy_document(document) + async with aclosing( + self.completer.get_completions_async(document2, complete_event) + ) as generator: + inner_completions = [completion async for completion in generator] + + for completion in self._fuzzy_completions( + word_before_cursor, inner_completions + ): + yield completion + else: + async with aclosing( + self.completer.get_completions_async(document, complete_event) + ) as generator: + async for completion in generator: + yield completion + + def _get_fuzzy_document(self, document: Document) -> tuple[str, Document]: word_before_cursor = document.get_word_before_cursor( pattern=re.compile(self._get_pattern()) ) - - # Get completions document2 = Document( text=document.text[: document.cursor_position - len(word_before_cursor)], cursor_position=document.cursor_position - len(word_before_cursor), ) + return word_before_cursor, document2 - inner_completions = list( - self.completer.get_completions(document2, complete_event) + def _get_fuzzy_completions( + self, document: Document, complete_event: CompleteEvent + ) -> Iterable[Completion]: + word_before_cursor, document2 = self._get_fuzzy_document(document) + yield from self._fuzzy_completions( + word_before_cursor, + list(self.completer.get_completions(document2, complete_event)), ) + def _fuzzy_completions( + self, word_before_cursor: str, inner_completions: Iterable[Completion] + ) -> Iterable[Completion]: fuzzy_matches: list[_FuzzyMatch] = [] if word_before_cursor == "": diff --git a/tests/test_fuzzy_async.py b/tests/test_fuzzy_async.py new file mode 100644 index 000000000..3a3c4dc6b --- /dev/null +++ b/tests/test_fuzzy_async.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from prompt_toolkit.completion import ( + CompleteEvent, + Completer, + Completion, + FuzzyCompleter, + ThreadedCompleter, + WordCompleter, +) +from prompt_toolkit.document import Document + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_fuzzy_completer_uses_async_implementation(enabled): + documents = [] + event = CompleteEvent(completion_requested=True) + + class AsyncCompleter(Completer): + def get_completions(self, document, complete_event): + raise AssertionError("Synchronous implementation must not be called") + + async def get_completions_async(self, document, complete_event): + documents.append(document.text) + assert complete_event is event + for word in ["leopard", "cat", "oar", "dinosaur"]: + yield Completion(word) + + async def collect(): + return [ + c + async for c in FuzzyCompleter( + AsyncCompleter(), enable_fuzzy=enabled + ).get_completions_async(Document("oar"), event) + ] + + completions = asyncio.run(collect()) + assert documents == ([""] if enabled else ["oar"]) + assert [c.text for c in completions] == ( + ["oar", "leopard", "dinosaur"] + if enabled + else ["leopard", "cat", "oar", "dinosaur"] + ) + assert [c.start_position for c in completions] == ([-3] * 3 if enabled else [0] * 4) + + +def test_fuzzy_completer_keeps_threaded_completions_off_event_loop(): + main_thread = threading.get_ident() + threads = [] + + class RecordingCompleter(WordCompleter): + def get_completions(self, document, complete_event): + threads.append(threading.get_ident()) + return super().get_completions(document, complete_event) + + async def collect(): + completer = FuzzyCompleter(ThreadedCompleter(RecordingCompleter(["leopard"]))) + return [ + c.text + async for c in completer.get_completions_async( + Document("oar"), CompleteEvent() + ) + ] + + assert asyncio.run(collect()) == ["leopard"] + assert threads and all(thread != main_thread for thread in threads) + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_fuzzy_completer_closes_async_generator(enabled): + closed = [] + + class AsyncCompleter(Completer): + def get_completions(self, document, complete_event): + raise AssertionError("Synchronous implementation must not be called") + + async def get_completions_async(self, document, complete_event): + try: + yield Completion("leopard") + if enabled: + raise ValueError("Completion failed") + yield Completion("dinosaur") + finally: + closed.append(True) + + async def consume(): + generator = FuzzyCompleter( + AsyncCompleter(), enable_fuzzy=enabled + ).get_completions_async(Document("oar"), CompleteEvent()) + if enabled: + with pytest.raises(ValueError, match="Completion failed"): + await anext(generator) + else: + assert (await anext(generator)).text == "leopard" + await generator.aclose() + assert closed == [True] + + asyncio.run(consume()) + + +@pytest.mark.parametrize("text", ["", "oar", "animals oar", "zz", "LEO"]) +def test_fuzzy_completer_async_matches_sync(text): + completer = FuzzyCompleter(WordCompleter(["leopard", "cat", "oar", "dinosaur"])) + document = Document(text) + event = CompleteEvent() + + async def collect(): + return [c async for c in completer.get_completions_async(document, event)] + + def describe(completions): + return [(c.text, c.start_position, c.display) for c in completions] + + assert describe(asyncio.run(collect())) == describe( + completer.get_completions(document, event) + ) + + +def test_fuzzy_completer_cancellation_closes_pending_generator(): + closed = [] + + async def consume(): + started = asyncio.Event() + + class WaitingCompleter(Completer): + def get_completions(self, document, complete_event): + raise AssertionError("Synchronous implementation must not be called") + + async def get_completions_async(self, document, complete_event): + try: + yield Completion("leopard") + started.set() + await asyncio.Event().wait() + finally: + closed.append(True) + + generator = FuzzyCompleter(WaitingCompleter()).get_completions_async( + Document("oar"), CompleteEvent() + ) + task = asyncio.create_task(anext(generator)) + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert closed == [True] + + asyncio.run(consume())