From 1e51d378075a54559c257d91ff2e516bfced6cbc 2023-01-08 18:16:23 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: 2023-01-08 18:16:23 Subject: [PATCH] Implement token-by-token autosuggestions --- diff --git a/IPython/terminal/shortcuts/__init__.py b/IPython/terminal/shortcuts/__init__.py index 68eaf65..f25a8c1 100644 --- a/IPython/terminal/shortcuts/__init__.py +++ b/IPython/terminal/shortcuts/__init__.py @@ -333,6 +333,7 @@ def create_ipython_shortcuts(shell, for_all_platforms: bool = False): kb.add("escape", "f", filter=focused_insert_vi & ebivim)( autosuggestions.accept_word ) + kb.add("c-right", filter=has_focus(DEFAULT_BUFFER))(autosuggestions.accept_token) # Simple Control keybindings key_cmd_dict = { diff --git a/IPython/terminal/shortcuts/autosuggestions.py b/IPython/terminal/shortcuts/autosuggestions.py index fe1e8d0..158d988 100644 --- a/IPython/terminal/shortcuts/autosuggestions.py +++ b/IPython/terminal/shortcuts/autosuggestions.py @@ -1,7 +1,13 @@ import re +import tokenize +from io import StringIO +from typing import List, Optional + from prompt_toolkit.key_binding import KeyPressEvent from prompt_toolkit.key_binding.bindings import named_commands as nc +from IPython.utils.tokenutil import generate_tokens + # Needed for to accept autosuggestions in vi insert mode def accept_in_vi_insert_mode(event: KeyPressEvent): @@ -18,7 +24,7 @@ def accept_in_vi_insert_mode(event: KeyPressEvent): nc.end_of_line(event) -def accept(event): +def accept(event: KeyPressEvent): """Accept suggestion""" b = event.current_buffer suggestion = b.suggestion @@ -28,7 +34,7 @@ def accept(event): nc.forward_char(event) -def accept_word(event): +def accept_word(event: KeyPressEvent): """Fill partial suggestion by word""" b = event.current_buffer suggestion = b.suggestion @@ -37,3 +43,45 @@ def accept_word(event): b.insert_text(next((x for x in t if x), "")) else: nc.forward_word(event) + + +def accept_token(event: KeyPressEvent): + """Fill partial suggestion by token""" + b = event.current_buffer + suggestion = b.suggestion + + if suggestion: + prefix = b.text + text = prefix + suggestion.text + + tokens: List[Optional[str]] = [None, None, None] + substings = [""] + i = 0 + + for token in generate_tokens(StringIO(text).readline): + if token.type == tokenize.NEWLINE: + index = len(text) + else: + index = text.index(token[1], len(substings[-1])) + substings.append(text[:index]) + tokenized_so_far = substings[-1] + if tokenized_so_far.startswith(prefix): + if i == 0 and len(tokenized_so_far) > len(prefix): + tokens[0] = tokenized_so_far[len(prefix) :] + substings.append(tokenized_so_far) + i += 1 + tokens[i] = token[1] + if i == 2: + break + i += 1 + + if tokens[0]: + to_insert: str + insert_text = substings[-2] + if tokens[1] and len(tokens[1]) == 1: + insert_text = substings[-1] + to_insert = insert_text[len(prefix) :] + b.insert_text(to_insert) + return + + nc.forward_word(event) diff --git a/IPython/tests/test_shortcuts.py b/IPython/terminal/tests/test_shortcuts.py similarity index 52% rename from IPython/tests/test_shortcuts.py rename to IPython/terminal/tests/test_shortcuts.py index 42edb92..92242f7 100644 --- a/IPython/tests/test_shortcuts.py +++ b/IPython/terminal/tests/test_shortcuts.py @@ -1,13 +1,17 @@ import pytest -from IPython.terminal.shortcuts import _apply_autosuggest +from IPython.terminal.shortcuts.autosuggestions import ( + accept_in_vi_insert_mode, + accept_token, +) -from unittest.mock import Mock +from unittest.mock import patch, Mock def make_event(text, cursor, suggestion): event = Mock() event.current_buffer = Mock() event.current_buffer.suggestion = Mock() + event.current_buffer.text = text event.current_buffer.cursor_position = cursor event.current_buffer.suggestion.text = suggestion event.current_buffer.document = Mock() @@ -32,9 +36,59 @@ def test_autosuggest_at_EOL(text, cursor, suggestion, called): event = make_event(text, cursor, suggestion) event.current_buffer.insert_text = Mock() - _apply_autosuggest(event) + accept_in_vi_insert_mode(event) if called: event.current_buffer.insert_text.assert_called() else: event.current_buffer.insert_text.assert_not_called() # event.current_buffer.document.get_end_of_line_position.assert_called() + + +@pytest.mark.parametrize( + "text, suggestion, expected", + [ + ("", "def out(tag: str, n=50):", "def "), + ("d", "ef out(tag: str, n=50):", "ef "), + ("de ", "f out(tag: str, n=50):", "f "), + ("def", " out(tag: str, n=50):", " "), + ("def ", "out(tag: str, n=50):", "out("), + ("def o", "ut(tag: str, n=50):", "ut("), + ("def ou", "t(tag: str, n=50):", "t("), + ("def out", "(tag: str, n=50):", "("), + ("def out(", "tag: str, n=50):", "tag: "), + ("def out(t", "ag: str, n=50):", "ag: "), + ("def out(ta", "g: str, n=50):", "g: "), + ("def out(tag", ": str, n=50):", ": "), + ("def out(tag:", " str, n=50):", " "), + ("def out(tag: ", "str, n=50):", "str, "), + ("def out(tag: s", "tr, n=50):", "tr, "), + ("def out(tag: st", "r, n=50):", "r, "), + ("def out(tag: str", ", n=50):", ", n"), + ("def out(tag: str,", " n=50):", " n"), + ("def out(tag: str, ", "n=50):", "n="), + ("def out(tag: str, n", "=50):", "="), + ("def out(tag: str, n=", "50):", "50)"), + ("def out(tag: str, n=5", "0):", "0)"), + ("def out(tag: str, n=50", "):", "):"), + ("def out(tag: str, n=50)", ":", ":"), + ], +) +def test_autosuggest_token(text, suggestion, expected): + event = make_event(text, len(text), suggestion) + event.current_buffer.insert_text = Mock() + accept_token(event) + assert event.current_buffer.insert_text.called + assert event.current_buffer.insert_text.call_args[0] == (expected,) + + +def test_autosuggest_token_empty(): + full = "def out(tag: str, n=50):" + event = make_event(full, len(full), "") + event.current_buffer.insert_text = Mock() + + with patch( + "prompt_toolkit.key_binding.bindings.named_commands.forward_word" + ) as forward_word: + accept_token(event) + assert not event.current_buffer.insert_text.called + assert forward_word.called