Preserve asynchronous calls in conditional validators and suggestions

src/prompt_toolkit/auto_suggest.py

@@ -156,6 +156,14 @@ def get_suggestion(self, buffer: Buffer, document: Document) -> Suggestion | Non
 
         return None
 
+    async def get_suggestion_async(
+        self, buff: Buffer, document: Document
+    ) -> Suggestion | None:
+        if self.filter():
+            return await self.auto_suggest.get_suggestion_async(buff, document)
+
+        return None
+
 
 class DynamicAutoSuggest(AutoSuggest):
     """

src/prompt_toolkit/validation.py

@@ -172,6 +172,10 @@ def validate(self, document: Document) -> None:
         if self.filter():
             self.validator.validate(document)
 
+    async def validate_async(self, document: Document) -> None:
+        if self.filter():
+            await self.validator.validate_async(document)
+
 
 class DynamicValidator(Validator):
     """

tests/test_auto_suggest.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from asyncio import run
+from threading import get_ident
+from unittest.mock import Mock
+
+import pytest
+
+from prompt_toolkit.auto_suggest import (
+    AutoSuggest,
+    ConditionalAutoSuggest,
+    Suggestion,
+    ThreadedAutoSuggest,
+)
+from prompt_toolkit.buffer import Buffer
+from prompt_toolkit.document import Document
+from prompt_toolkit.filters import Condition
+
+
+def test_conditional_auto_suggest_async():
+    auto_suggest = Mock(spec=AutoSuggest)
+    suggestion = Suggestion("completion")
+    auto_suggest.get_suggestion_async.return_value = suggestion
+    enabled = False
+    conditional = ConditionalAutoSuggest(auto_suggest, Condition(lambda: enabled))
+    buffer = Buffer()
+    document = Document("input")
+
+    async def suggest():
+        nonlocal enabled
+        assert await conditional.get_suggestion_async(buffer, document) is None
+        auto_suggest.get_suggestion_async.assert_not_called()
+
+        enabled = True
+        assert await conditional.get_suggestion_async(buffer, document) is suggestion
+        auto_suggest.get_suggestion_async.assert_awaited_once_with(buffer, document)
+        auto_suggest.get_suggestion.assert_not_called()
+
+        enabled = False
+        assert await conditional.get_suggestion_async(buffer, document) is None
+        auto_suggest.get_suggestion_async.assert_awaited_once_with(buffer, document)
+
+    run(suggest())
+
+
+def test_conditional_auto_suggest_async_error():
+    auto_suggest = Mock(spec=AutoSuggest)
+    error = RuntimeError("Suggestion failed")
+    auto_suggest.get_suggestion_async.side_effect = error
+    conditional = ConditionalAutoSuggest(auto_suggest, True)
+
+    with pytest.raises(RuntimeError) as exc_info:
+        run(conditional.get_suggestion_async(Buffer(), Document("input")))
+
+    assert exc_info.value is error
+    auto_suggest.get_suggestion.assert_not_called()
+
+
+def test_conditional_threaded_auto_suggest():
+    thread_ids = []
+    suggestion = Suggestion("completion")
+
+    class TestAutoSuggest(AutoSuggest):
+        def get_suggestion(self, buffer, document):
+            thread_ids.append(get_ident())
+            return suggestion
+
+    conditional = ConditionalAutoSuggest(ThreadedAutoSuggest(TestAutoSuggest()), True)
+    buffer = Buffer()
+    document = Document("input")
+
+    assert conditional.get_suggestion(buffer, document) is suggestion
+    assert run(conditional.get_suggestion_async(buffer, document)) is suggestion
+    assert len(thread_ids) == 2
+    assert thread_ids[0] == get_ident()
+    assert thread_ids[1] != get_ident()
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_conditional_sync_auto_suggest_async(enabled):
+    calls = []
+    suggestion = Suggestion("completion")
+
+    class TestAutoSuggest(AutoSuggest):
+        def get_suggestion(self, buffer, document):
+            calls.append((buffer, document))
+            return suggestion
+
+    conditional = ConditionalAutoSuggest(TestAutoSuggest(), enabled)
+    buffer = Buffer()
+    document = Document("input")
+
+    result = run(conditional.get_suggestion_async(buffer, document))
+
+    assert result is (suggestion if enabled else None)
+    assert calls == ([(buffer, document)] if enabled else [])

tests/test_validation.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from asyncio import run
+from threading import get_ident
+from unittest.mock import Mock
+
+import pytest
+
+from prompt_toolkit.document import Document
+from prompt_toolkit.filters import Condition
+from prompt_toolkit.validation import (
+    ConditionalValidator,
+    ThreadedValidator,
+    ValidationError,
+    Validator,
+)
+
+
+def test_conditional_validator_async():
+    validator = Mock(spec=Validator)
+    enabled = False
+    conditional = ConditionalValidator(validator, Condition(lambda: enabled))
+    document = Document("input")
+
+    async def validate():
+        nonlocal enabled
+        await conditional.validate_async(document)
+        validator.validate_async.assert_not_called()
+
+        enabled = True
+        await conditional.validate_async(document)
+        validator.validate_async.assert_awaited_once_with(document)
+        validator.validate.assert_not_called()
+
+        enabled = False
+        await conditional.validate_async(document)
+        validator.validate_async.assert_awaited_once_with(document)
+
+    run(validate())
+
+
+def test_conditional_validator_async_error():
+    validator = Mock(spec=Validator)
+    error = ValidationError(cursor_position=2, message="Invalid input")
+    validator.validate_async.side_effect = error
+    conditional = ConditionalValidator(validator, True)
+
+    with pytest.raises(ValidationError) as exc_info:
+        run(conditional.validate_async(Document("input")))
+
+    assert exc_info.value is error
+    validator.validate.assert_not_called()
+
+
+def test_conditional_threaded_validator():
+    thread_ids = []
+    validator = Validator.from_callable(
+        lambda text: thread_ids.append(get_ident()) is None
+    )
+    conditional = ConditionalValidator(ThreadedValidator(validator), True)
+    document = Document("input")
+
+    conditional.validate(document)
+    run(conditional.validate_async(document))
+
+    assert len(thread_ids) == 2
+    assert thread_ids[0] == get_ident()
+    assert thread_ids[1] != get_ident()
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_conditional_sync_validator_async(enabled):
+    calls = []
+    validator = Validator.from_callable(lambda text: calls.append(text) is None)
+    conditional = ConditionalValidator(validator, enabled)
+
+    run(conditional.validate_async(Document("input")))
+
+    assert calls == (["input"] if enabled else [])