1  import pytest
   2  from IPython.terminal.shortcuts.auto_suggest import (
   3      accept,
   4      accept_or_jump_to_end,
   5      accept_token,
   6      accept_character,
   7      accept_word,
   8      accept_and_keep_cursor,
   9      discard,
  10      llm_autosuggestion,
  11      NavigableAutoSuggestFromHistory,
  12      swap_autosuggestion_up,
  13      swap_autosuggestion_down,
  14  )
  15  from IPython.terminal.shortcuts.auto_match import skip_over
  16  from IPython.terminal.shortcuts import create_ipython_shortcuts, reset_search_buffer
  17  from IPython.testing import decorators as dec
  18  
  19  from prompt_toolkit.history import InMemoryHistory
  20  from prompt_toolkit.buffer import Buffer
  21  from prompt_toolkit.document import Document
  22  from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
  23  from prompt_toolkit.enums import DEFAULT_BUFFER
  24  
  25  from unittest.mock import patch, Mock
  26  
  27  
  28  def test_deprected():
  29      import IPython.terminal.shortcuts.auto_suggest as iptsa
  30  
  31      with pytest.warns(DeprecationWarning, match=r"8\.12.+accept_or_jump_to_end"):
  32          iptsa.accept_in_vi_insert_mode
  33  
  34  
  35  def make_event(text, cursor, suggestion):
  36      event = Mock()
  37      event.current_buffer = Mock()
  38      event.current_buffer.suggestion = Mock()
  39      event.current_buffer.text = text
  40      event.current_buffer.cursor_position = cursor
  41      event.current_buffer.suggestion.text = suggestion
  42      event.current_buffer.document = Document(text=text, cursor_position=cursor)
  43      return event
  44  
  45  
  46  try:
  47      from .fake_llm import FIBONACCI
  48  except ImportError:
  49      FIBONACCI = None
  50  
  51  
  52  @dec.skip_without("jupyter_ai")
  53  @pytest.mark.asyncio
  54  async def test_llm_autosuggestion():
  55      provider = NavigableAutoSuggestFromHistory()
  56      ip = get_ipython()
  57      ip.auto_suggest = provider
  58      ip.llm_provider_class = "IPython.terminal.tests.fake_llm.FibonacciCompletionProvider"
  59      text = "def fib"
  60      event = make_event(text, len(text), "")
  61      event.current_buffer.history.shell.history_manager.get_range = Mock(return_value=[])
  62      await llm_autosuggestion(event)
  63      assert event.current_buffer.suggestion.text == FIBONACCI[len(text) :]
  64  
  65  
  66  @pytest.mark.parametrize(
  67      "text, suggestion, expected",
  68      [
  69          ("", "def out(tag: str, n=50):", "def out(tag: str, n=50):"),
  70          ("def ", "out(tag: str, n=50):", "out(tag: str, n=50):"),
  71      ],
  72  )
  73  def test_accept(text, suggestion, expected):
  74      event = make_event(text, len(text), suggestion)
  75      buffer = event.current_buffer
  76      buffer.insert_text = Mock()
  77      accept(event)
  78      assert buffer.insert_text.called
  79      assert buffer.insert_text.call_args[0] == (expected,)
  80  
  81  
  82  @pytest.mark.parametrize(
  83      "text, suggestion",
  84      [
  85          ("", "def out(tag: str, n=50):"),
  86          ("def ", "out(tag: str, n=50):"),
  87      ],
  88  )
  89  def test_discard(text, suggestion):
  90      event = make_event(text, len(text), suggestion)
  91      buffer = event.current_buffer
  92      buffer.insert_text = Mock()
  93      discard(event)
  94      assert not buffer.insert_text.called
  95      assert buffer.suggestion is None
  96  
  97  
  98  @pytest.mark.parametrize(
  99      "text, cursor, suggestion, called",
 100      [
 101          ("123456", 6, "123456789", True),
 102          ("123456", 3, "123456789", False),
 103          ("123456   \n789", 6, "123456789", True),
 104      ],
 105  )
 106  def test_autosuggest_at_EOL(text, cursor, suggestion, called):
 107      """
 108      test that autosuggest is only applied at end of line.
 109      """
 110  
 111      event = make_event(text, cursor, suggestion)
 112      event.current_buffer.insert_text = Mock()
 113      accept_or_jump_to_end(event)
 114      if called:
 115          event.current_buffer.insert_text.assert_called()
 116      else:
 117          event.current_buffer.insert_text.assert_not_called()
 118          # event.current_buffer.document.get_end_of_line_position.assert_called()
 119  
 120  
 121  @pytest.mark.parametrize(
 122      "text, suggestion, expected",
 123      [
 124          ("", "def out(tag: str, n=50):", "def "),
 125          ("d", "ef out(tag: str, n=50):", "ef "),
 126          ("de ", "f out(tag: str, n=50):", "f "),
 127          ("def", " out(tag: str, n=50):", " "),
 128          ("def ", "out(tag: str, n=50):", "out("),
 129          ("def o", "ut(tag: str, n=50):", "ut("),
 130          ("def ou", "t(tag: str, n=50):", "t("),
 131          ("def out", "(tag: str, n=50):", "("),
 132          ("def out(", "tag: str, n=50):", "tag: "),
 133          ("def out(t", "ag: str, n=50):", "ag: "),
 134          ("def out(ta", "g: str, n=50):", "g: "),
 135          ("def out(tag", ": str, n=50):", ": "),
 136          ("def out(tag:", " str, n=50):", " "),
 137          ("def out(tag: ", "str, n=50):", "str, "),
 138          ("def out(tag: s", "tr, n=50):", "tr, "),
 139          ("def out(tag: st", "r, n=50):", "r, "),
 140          ("def out(tag: str", ", n=50):", ", n"),
 141          ("def out(tag: str,", " n=50):", " n"),
 142          ("def out(tag: str, ", "n=50):", "n="),
 143          ("def out(tag: str, n", "=50):", "="),
 144          ("def out(tag: str, n=", "50):", "50)"),
 145          ("def out(tag: str, n=5", "0):", "0)"),
 146          ("def out(tag: str, n=50", "):", "):"),
 147          ("def out(tag: str, n=50)", ":", ":"),
 148      ],
 149  )
 150  def test_autosuggest_token(text, suggestion, expected):
 151      event = make_event(text, len(text), suggestion)
 152      event.current_buffer.insert_text = Mock()
 153      accept_token(event)
 154      assert event.current_buffer.insert_text.called
 155      assert event.current_buffer.insert_text.call_args[0] == (expected,)
 156  
 157  
 158  @pytest.mark.parametrize(
 159      "text, suggestion, expected",
 160      [
 161          ("", "def out(tag: str, n=50):", "d"),
 162          ("d", "ef out(tag: str, n=50):", "e"),
 163          ("de ", "f out(tag: str, n=50):", "f"),
 164          ("def", " out(tag: str, n=50):", " "),
 165      ],
 166  )
 167  def test_accept_character(text, suggestion, expected):
 168      event = make_event(text, len(text), suggestion)
 169      event.current_buffer.insert_text = Mock()
 170      accept_character(event)
 171      assert event.current_buffer.insert_text.called
 172      assert event.current_buffer.insert_text.call_args[0] == (expected,)
 173  
 174  
 175  @pytest.mark.parametrize(
 176      "text, suggestion, expected",
 177      [
 178          ("", "def out(tag: str, n=50):", "def "),
 179          ("d", "ef out(tag: str, n=50):", "ef "),
 180          ("de", "f out(tag: str, n=50):", "f "),
 181          ("def", " out(tag: str, n=50):", " "),
 182          # (this is why we also have accept_token)
 183          ("def ", "out(tag: str, n=50):", "out(tag: "),
 184      ],
 185  )
 186  def test_accept_word(text, suggestion, expected):
 187      event = make_event(text, len(text), suggestion)
 188      event.current_buffer.insert_text = Mock()
 189      accept_word(event)
 190      assert event.current_buffer.insert_text.called
 191      assert event.current_buffer.insert_text.call_args[0] == (expected,)
 192  
 193  
 194  @pytest.mark.parametrize(
 195      "text, suggestion, expected, cursor",
 196      [
 197          ("", "def out(tag: str, n=50):", "def out(tag: str, n=50):", 0),
 198          ("def ", "out(tag: str, n=50):", "out(tag: str, n=50):", 4),
 199      ],
 200  )
 201  def test_accept_and_keep_cursor(text, suggestion, expected, cursor):
 202      event = make_event(text, cursor, suggestion)
 203      buffer = event.current_buffer
 204      buffer.insert_text = Mock()
 205      accept_and_keep_cursor(event)
 206      assert buffer.insert_text.called
 207      assert buffer.insert_text.call_args[0] == (expected,)
 208      assert buffer.cursor_position == cursor
 209  
 210  
 211  def test_autosuggest_token_empty():
 212      full = "def out(tag: str, n=50):"
 213      event = make_event(full, len(full), "")
 214      event.current_buffer.insert_text = Mock()
 215  
 216      with patch(
 217          "prompt_toolkit.key_binding.bindings.named_commands.forward_word"
 218      ) as forward_word:
 219          accept_token(event)
 220          assert not event.current_buffer.insert_text.called
 221          assert forward_word.called
 222  
 223  
 224  def test_reset_search_buffer():
 225      event_with_text = Mock()
 226      event_with_text.current_buffer.document.text = "some text"
 227      event_with_text.current_buffer.reset = Mock()
 228  
 229      event_empty = Mock()
 230      event_empty.current_buffer.document.text = ""
 231      event_empty.app.layout.focus = Mock()
 232  
 233      reset_search_buffer(event_with_text)
 234      event_with_text.current_buffer.reset.assert_called_once()
 235  
 236      reset_search_buffer(event_empty)
 237      event_empty.app.layout.focus.assert_called_once_with(DEFAULT_BUFFER)
 238  
 239  
 240  def test_other_providers():
 241      """Ensure that swapping autosuggestions does not break with other providers"""
 242      provider = AutoSuggestFromHistory()
 243      ip = get_ipython()
 244      ip.auto_suggest = provider
 245      event = Mock()
 246      event.current_buffer = Buffer()
 247      assert swap_autosuggestion_up(event) is None
 248      assert swap_autosuggestion_down(event) is None
 249  
 250  
 251  @pytest.mark.asyncio
 252  async def test_navigable_provider():
 253      provider = NavigableAutoSuggestFromHistory()
 254      history = InMemoryHistory(history_strings=["very_a", "very", "very_b", "very_c"])
 255      buffer = Buffer(history=history)
 256      ip = get_ipython()
 257      ip.auto_suggest = provider
 258  
 259      async for _ in history.load():
 260          pass
 261  
 262      buffer.cursor_position = 5
 263      buffer.text = "very"
 264  
 265      up = swap_autosuggestion_up
 266      down = swap_autosuggestion_down
 267  
 268      event = Mock()
 269      event.current_buffer = buffer
 270  
 271      def get_suggestion():
 272          suggestion = provider.get_suggestion(buffer, buffer.document)
 273          buffer.suggestion = suggestion
 274          return suggestion
 275  
 276      assert get_suggestion().text == "_c"
 277  
 278      # should go up
 279      up(event)
 280      assert get_suggestion().text == "_b"
 281  
 282      # should skip over 'very' which is identical to buffer content
 283      up(event)
 284      assert get_suggestion().text == "_a"
 285  
 286      # should cycle back to beginning
 287      up(event)
 288      assert get_suggestion().text == "_c"
 289  
 290      # should cycle back through end boundary
 291      down(event)
 292      assert get_suggestion().text == "_a"
 293  
 294      down(event)
 295      assert get_suggestion().text == "_b"
 296  
 297      down(event)
 298      assert get_suggestion().text == "_c"
 299  
 300      down(event)
 301      assert get_suggestion().text == "_a"
 302  
 303  
 304  @pytest.mark.asyncio
 305  async def test_navigable_provider_multiline_entries():
 306      provider = NavigableAutoSuggestFromHistory()
 307      history = InMemoryHistory(history_strings=["very_a\nvery_b", "very_c"])
 308      buffer = Buffer(history=history)
 309      ip = get_ipython()
 310      ip.auto_suggest = provider
 311  
 312      async for _ in history.load():
 313          pass
 314  
 315      buffer.cursor_position = 5
 316      buffer.text = "very"
 317      up = swap_autosuggestion_up
 318      down = swap_autosuggestion_down
 319  
 320      event = Mock()
 321      event.current_buffer = buffer
 322  
 323      def get_suggestion():
 324          suggestion = provider.get_suggestion(buffer, buffer.document)
 325          buffer.suggestion = suggestion
 326          return suggestion
 327  
 328      assert get_suggestion().text == "_c"
 329  
 330      up(event)
 331      assert get_suggestion().text == "_b"
 332  
 333      up(event)
 334      assert get_suggestion().text == "_a"
 335  
 336      down(event)
 337      assert get_suggestion().text == "_b"
 338  
 339      down(event)
 340      assert get_suggestion().text == "_c"
 341  
 342  
 343  def create_session_mock():
 344      session = Mock()
 345      session.default_buffer = Buffer()
 346      return session
 347  
 348  
 349  def test_navigable_provider_connection():
 350      provider = NavigableAutoSuggestFromHistory()
 351      provider.skip_lines = 1
 352  
 353      session_1 = create_session_mock()
 354      provider.connect(session_1)
 355  
 356      assert provider.skip_lines == 1
 357      session_1.default_buffer.on_text_insert.fire()
 358      assert provider.skip_lines == 0
 359  
 360      session_2 = create_session_mock()
 361      provider.connect(session_2)
 362      provider.skip_lines = 2
 363  
 364      assert provider.skip_lines == 2
 365      session_2.default_buffer.on_text_insert.fire()
 366      assert provider.skip_lines == 0
 367  
 368      provider.skip_lines = 3
 369      provider.disconnect()
 370      session_1.default_buffer.on_text_insert.fire()
 371      session_2.default_buffer.on_text_insert.fire()
 372      assert provider.skip_lines == 3
 373  
 374  
 375  @pytest.fixture
 376  def ipython_with_prompt():
 377      ip = get_ipython()
 378      ip.pt_app = Mock()
 379      ip.pt_app.key_bindings = create_ipython_shortcuts(ip)
 380      try:
 381          yield ip
 382      finally:
 383          ip.pt_app = None
 384  
 385  
 386  def find_bindings_by_command(command):
 387      ip = get_ipython()
 388      return [
 389          binding
 390          for binding in ip.pt_app.key_bindings.bindings
 391          if binding.handler == command
 392      ]
 393  
 394  
 395  def test_modify_unique_shortcut(ipython_with_prompt):
 396      original = find_bindings_by_command(accept_token)
 397      assert len(original) == 1
 398  
 399      ipython_with_prompt.shortcuts = [
 400          {"command": "IPython:auto_suggest.accept_token", "new_keys": ["a", "b", "c"]}
 401      ]
 402      matched = find_bindings_by_command(accept_token)
 403      assert len(matched) == 1
 404      assert list(matched[0].keys) == ["a", "b", "c"]
 405      assert list(matched[0].keys) != list(original[0].keys)
 406      assert matched[0].filter == original[0].filter
 407  
 408      ipython_with_prompt.shortcuts = [
 409          {"command": "IPython:auto_suggest.accept_token", "new_filter": "always"}
 410      ]
 411      matched = find_bindings_by_command(accept_token)
 412      assert len(matched) == 1
 413      assert list(matched[0].keys) != ["a", "b", "c"]
 414      assert list(matched[0].keys) == list(original[0].keys)
 415      assert matched[0].filter != original[0].filter
 416  
 417  
 418  def test_disable_shortcut(ipython_with_prompt):
 419      matched = find_bindings_by_command(accept_token)
 420      assert len(matched) == 1
 421  
 422      ipython_with_prompt.shortcuts = [
 423          {"command": "IPython:auto_suggest.accept_token", "new_keys": []}
 424      ]
 425      matched = find_bindings_by_command(accept_token)
 426      assert len(matched) == 0
 427  
 428      ipython_with_prompt.shortcuts = []
 429      matched = find_bindings_by_command(accept_token)
 430      assert len(matched) == 1
 431  
 432  
 433  def test_modify_shortcut_with_filters(ipython_with_prompt):
 434      matched = find_bindings_by_command(skip_over)
 435      matched_keys = {m.keys[0] for m in matched}
 436      assert matched_keys == {")", "]", "}", "'", '"'}
 437  
 438      with pytest.raises(ValueError, match="Multiple shortcuts matching"):
 439          ipython_with_prompt.shortcuts = [
 440              {"command": "IPython:auto_match.skip_over", "new_keys": ["x"]}
 441          ]
 442  
 443      ipython_with_prompt.shortcuts = [
 444          {
 445              "command": "IPython:auto_match.skip_over",
 446              "new_keys": ["x"],
 447              "match_filter": "focused_insert & auto_match & followed_by_single_quote",
 448          }
 449      ]
 450      matched = find_bindings_by_command(skip_over)
 451      matched_keys = {m.keys[0] for m in matched}
 452      assert matched_keys == {")", "]", "}", "x", '"'}
 453  
 454  
 455  def example_command():
 456      pass
 457  
 458  
 459  def test_add_shortcut_for_new_command(ipython_with_prompt):
 460      matched = find_bindings_by_command(example_command)
 461      assert len(matched) == 0
 462  
 463      with pytest.raises(ValueError, match="example_command is not a known"):
 464          ipython_with_prompt.shortcuts = [
 465              {"command": "example_command", "new_keys": ["x"]}
 466          ]
 467      matched = find_bindings_by_command(example_command)
 468      assert len(matched) == 0
 469  
 470  
 471  def test_modify_shortcut_failure(ipython_with_prompt):
 472      with pytest.raises(ValueError, match="No shortcuts matching"):
 473          ipython_with_prompt.shortcuts = [
 474              {
 475                  "command": "IPython:auto_match.skip_over",
 476                  "match_keys": ["x"],
 477                  "new_keys": ["y"],
 478              }
 479          ]
 480  
 481  
 482  def test_add_shortcut_for_existing_command(ipython_with_prompt):
 483      matched = find_bindings_by_command(skip_over)
 484      assert len(matched) == 5
 485  
 486      with pytest.raises(ValueError, match="Cannot add a shortcut without keys"):
 487          ipython_with_prompt.shortcuts = [
 488              {"command": "IPython:auto_match.skip_over", "new_keys": [], "create": True}
 489          ]
 490  
 491      ipython_with_prompt.shortcuts = [
 492          {"command": "IPython:auto_match.skip_over", "new_keys": ["x"], "create": True}
 493      ]
 494      matched = find_bindings_by_command(skip_over)
 495      assert len(matched) == 6
 496  
 497      ipython_with_prompt.shortcuts = []
 498      matched = find_bindings_by_command(skip_over)
 499      assert len(matched) == 5
 500  
 501  
 502  def test_setting_shortcuts_before_pt_app_init():
 503      ipython = get_ipython()
 504      assert ipython.pt_app is None
 505      shortcuts = [
 506          {"command": "IPython:auto_match.skip_over", "new_keys": ["x"], "create": True}
 507      ]
 508      ipython.shortcuts = shortcuts
 509      assert ipython.shortcuts == shortcuts
 510