diff --git a/src/prompt_toolkit/key_binding/bindings/auto_suggest.py b/src/prompt_toolkit/key_binding/bindings/auto_suggest.py index b487f14ff..6f6842a5b 100644 --- a/src/prompt_toolkit/key_binding/bindings/auto_suggest.py +++ b/src/prompt_toolkit/key_binding/bindings/auto_suggest.py @@ -7,7 +7,7 @@ import re from prompt_toolkit.application.current import get_app -from prompt_toolkit.filters import Condition, emacs_mode +from prompt_toolkit.filters import Condition, emacs_mode, vi_navigation_mode from prompt_toolkit.key_binding.key_bindings import KeyBindings from prompt_toolkit.key_binding.key_processor import KeyPressEvent @@ -18,6 +18,26 @@ E = KeyPressEvent +@Condition +def suggestion_available() -> bool: + app = get_app() + buffer = app.current_buffer + document = buffer.document + + return ( + buffer.suggestion is not None + and len(buffer.suggestion.text) > 0 + and ( + document.is_cursor_at_the_end + or ( + vi_navigation_mode() + and document.cursor_position_row == document.line_count - 1 + and document.cursor_position_col == len(document.current_line) - 1 + ) + ) + ) + + def load_auto_suggest_bindings() -> KeyBindings: """ Key bindings for accepting auto suggestion text. @@ -29,15 +49,7 @@ def load_auto_suggest_bindings() -> KeyBindings: key_bindings = KeyBindings() handle = key_bindings.add - @Condition - def suggestion_available() -> bool: - app = get_app() - return ( - app.current_buffer.suggestion is not None - and len(app.current_buffer.suggestion.text) > 0 - and app.current_buffer.document.is_cursor_at_the_end - ) - + @handle("l", filter=suggestion_available & vi_navigation_mode) @handle("c-f", filter=suggestion_available) @handle("c-e", filter=suggestion_available) @handle("right", filter=suggestion_available) @@ -49,6 +61,8 @@ def _accept(event: E) -> None: suggestion = b.suggestion if suggestion: + if vi_navigation_mode(): + b.cursor_position = len(b.text) b.insert_text(suggestion.text) @handle("escape", "f", filter=suggestion_available & emacs_mode) diff --git a/src/prompt_toolkit/layout/processors.py b/src/prompt_toolkit/layout/processors.py index 139f07c3e..f12f7b5e9 100644 --- a/src/prompt_toolkit/layout/processors.py +++ b/src/prompt_toolkit/layout/processors.py @@ -16,7 +16,12 @@ from prompt_toolkit.application.current import get_app from prompt_toolkit.cache import SimpleCache from prompt_toolkit.document import Document -from prompt_toolkit.filters import FilterOrBool, to_filter, vi_insert_multiple_mode +from prompt_toolkit.filters import ( + FilterOrBool, + to_filter, + vi_insert_multiple_mode, + vi_navigation_mode, +) from prompt_toolkit.formatted_text import ( AnyFormattedText, StyleAndTextTuples, @@ -591,7 +596,13 @@ def apply_transformation(self, ti: TransformationInput) -> Transformation: if ti.lineno == ti.document.line_count - 1: buffer = ti.buffer_control.buffer - if buffer.suggestion and ti.document.is_cursor_at_the_end: + at_end = ti.document.is_cursor_at_the_end or ( + vi_navigation_mode() + and ti.document.cursor_position_row == ti.document.line_count - 1 + and ti.document.cursor_position_col == len(ti.document.current_line) - 1 + ) + + if buffer.suggestion and at_end: suggestion = buffer.suggestion.text else: suggestion = "" diff --git a/tests/test_cli.py b/tests/test_cli.py index a876f2993..1c6980368 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -9,6 +9,7 @@ import pytest +from prompt_toolkit.auto_suggest import Suggestion from prompt_toolkit.clipboard import ClipboardData, InMemoryClipboard from prompt_toolkit.enums import EditingMode from prompt_toolkit.filters import ViInsertMode @@ -17,6 +18,8 @@ from prompt_toolkit.input.vt100_parser import ANSI_SEQUENCES from prompt_toolkit.key_binding.bindings.named_commands import prefix_meta from prompt_toolkit.key_binding.key_bindings import KeyBindings +from prompt_toolkit.key_binding.vi_state import InputMode +from prompt_toolkit.layout.processors import AppendAutoSuggestion, TransformationInput from prompt_toolkit.output import DummyOutput from prompt_toolkit.shortcuts import PromptSession @@ -893,6 +896,41 @@ def test_vi_temp_navigation_mode(): assert result.cursor_position == 5 +@pytest.mark.parametrize("accept_key", ["l", "\x1b[C"]) +def test_vi_navigation_keeps_and_accepts_suggestion(accept_key): + with create_pipe_input() as inp: + inp.send_text(accept_key + "\r") + session = PromptSession( + input=inp, + output=DummyOutput(), + editing_mode=EditingMode.VI, + ) + + displayed = [] + + def pre_run(): + buffer = session.default_buffer + buffer.suggestion = Suggestion(" hello") + buffer.cursor_position = len(buffer.text) - 1 + session.app.vi_state.input_mode = InputMode.NAVIGATION + + transformation = AppendAutoSuggestion().apply_transformation( + TransformationInput( + session.app.layout.current_control, + buffer.document, + 0, + lambda i: i, + [("", buffer.text)], + 80, + 24, + ) + ) + displayed.extend(text for _, text in transformation.fragments) + + assert session.prompt(default="echo", pre_run=pre_run) == "echo hello" + assert "".join(displayed) == "echo hello" + + def test_vi_macros(): feed = partial(_feed_cli_with_input, editing_mode=EditingMode.VI)