Make search_state dynamic for key bindings. (Added property get_search_state to KeyBindingManager.)

3 files changed, 45 insertions, 23 deletions

diff --git a/prompt_toolkit/key_binding/bindings/emacs.py b/prompt_toolkit/key_binding/bindings/emacs.py
index 21db2590..e92abfeb 100644
--- a/prompt_toolkit/key_binding/bindings/emacs.py
+++ b/prompt_toolkit/key_binding/bindings/emacs.py
@@ -476,10 +476,15 @@ def load_emacs_system_bindings(registry, filter=None):
         event.cli.focus_stack.pop()
 
 
-def load_emacs_search_bindings(registry, filter=None):
+def load_emacs_search_bindings(registry, get_search_state=None, filter=None):
     handle = create_handle_decorator(registry, filter)
     has_focus = filters.HasFocus(SEARCH_BUFFER)
 
+    assert get_search_state is None or callable(get_search_state)
+
+    if not get_search_state:
+        def get_search_state(cli): return cli.search_state
+
     @handle(Keys.ControlG, filter=has_focus)
     @handle(Keys.ControlC, filter=has_focus)
     # NOTE: the reason for not also binding Escape to this one, is that we want
@@ -504,10 +509,10 @@ def load_emacs_search_bindings(registry, filter=None):
 
         # Update search state.
         if search_buffer.text:
-            event.cli.search_state.text = search_buffer.text
+            get_search_state(event.cli).text = search_buffer.text
 
         # Apply search.
-        input_buffer.apply_search(event.cli.search_state, include_current_position=True)
+        input_buffer.apply_search(get_search_state(event.cli), include_current_position=True)
 
         # Add query to history of search line.
         search_buffer.append_to_history()
@@ -518,19 +523,19 @@ def load_emacs_search_bindings(registry, filter=None):
 
     @handle(Keys.ControlR, filter= ~has_focus)
     def _(event):
-        event.cli.search_state.direction = IncrementalSearchDirection.BACKWARD
+        get_search_state(event.cli).direction = IncrementalSearchDirection.BACKWARD
         event.cli.focus_stack.push(SEARCH_BUFFER)
 
     @handle(Keys.ControlS, filter= ~has_focus)
     def _(event):
-        event.cli.search_state.direction = IncrementalSearchDirection.FORWARD
+        get_search_state(event.cli).direction = IncrementalSearchDirection.FORWARD
         event.cli.focus_stack.push(SEARCH_BUFFER)
 
     @handle(Keys.ControlR, filter=has_focus)
     @handle(Keys.Up, filter=has_focus)
     def _(event):
         # Update search_state.
-        search_state = event.cli.search_state
+        search_state = get_search_state(event.cli)
         direction_changed = search_state.direction != IncrementalSearchDirection.BACKWARD
 
         search_state.text = event.cli.buffers[SEARCH_BUFFER].text
@@ -539,14 +544,14 @@ def load_emacs_search_bindings(registry, filter=None):
         # Apply search to current buffer.
         if not direction_changed:
             input_buffer = event.cli.buffers[event.cli.focus_stack.previous]
-            input_buffer.apply_search(event.cli.search_state,
+            input_buffer.apply_search(search_state,
                                       include_current_position=False, count=event.arg)
 
     @handle(Keys.ControlS, filter=has_focus)
     @handle(Keys.Down, filter=has_focus)
     def _(event):
         # Update search_state.
-        search_state = event.cli.search_state
+        search_state = get_search_state(event.cli)
         direction_changed = search_state.direction != IncrementalSearchDirection.FORWARD
 
         search_state.text = event.cli.buffers[SEARCH_BUFFER].text
@@ -555,7 +560,7 @@ def load_emacs_search_bindings(registry, filter=None):
         # Apply search to current buffer.
         if not direction_changed:
             input_buffer = event.cli.buffers[event.cli.focus_stack.previous]
-            input_buffer.apply_search(event.cli.search_state,
+            input_buffer.apply_search(search_state,
                                       include_current_position=False, count=event.arg)
 
 
diff --git a/prompt_toolkit/key_binding/bindings/vi.py b/prompt_toolkit/key_binding/bindings/vi.py
index e476d4bd..31bec091 100644
--- a/prompt_toolkit/key_binding/bindings/vi.py
+++ b/prompt_toolkit/key_binding/bindings/vi.py
@@ -62,7 +62,7 @@ class CursorRegion(object):
             return self.end, self.start
 
 
-def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=None):
+def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), get_search_state=None, filter=None):
     """
     Vi extensions.
 
@@ -72,6 +72,7 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
     :param get_vi_state: Callable that takes a CommandLineInterface instances and returns the used ViState.
     :param enable_visual_key: Filter to enable lowercase 'v' bindings. A reason to disable these
          are to support open-in-editor functionality. These key bindings conflict.
+     :param get_search_state: None or a callable that takes a CommandLineInterface and returns a SearchState.
     """
     # Note: Some key bindings have the "~IsReadOnly()" filter added. This
     #       prevents the handler to be executed when the focus is on a
@@ -85,6 +86,10 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
     assert callable(get_vi_state)
     enable_visual_key = to_cli_filter(enable_visual_key)
 
+    # Default get_search_state.
+    if get_search_state is None:
+        def get_search_state(cli): return cli.search_state
+
     handle = create_handle_decorator(registry, filter)
 
     insert_mode = ViStateFilter(get_vi_state, InputMode.INSERT) & ~ filters.HasSelection()
@@ -344,7 +349,7 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
         Search next.
         """
         event.current_buffer.apply_search(
-            event.cli.search_state, include_current_position=False,
+            get_search_state(event.cli), include_current_position=False,
             count=event.arg)
 
     @handle('N', filter=navigation_mode)
@@ -353,7 +358,7 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
         Search previous.
         """
         event.current_buffer.apply_search(
-            ~event.cli.search_state, include_current_position=False,
+            ~get_search_state(event.cli), include_current_position=False,
             count=event.arg)
 
     @handle('p', filter=navigation_mode)
@@ -600,7 +605,7 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
         """
         b = event.cli.current_buffer
 
-        search_state = event.cli.search_state
+        search_state = get_search_state(event.cli)
         search_state.text = b.document.get_word_under_cursor()
         search_state.direction = IncrementalSearchDirection.BACKWARD
 
@@ -614,7 +619,7 @@ def load_vi_bindings(registry, get_vi_state, enable_visual_key=Always(), filter=
         """
         b = event.cli.current_buffer
 
-        search_state = event.cli.search_state
+        search_state = get_search_state(event.cli)
         search_state.text = b.document.get_word_under_cursor()
         search_state.direction = IncrementalSearchDirection.FORWARD
 
@@ -1259,8 +1264,13 @@ def load_vi_system_bindings(registry, get_vi_state, filter=None):
         event.cli.focus_stack.pop()
 
 
-def load_vi_search_bindings(registry, get_vi_state, filter=None, search_buffer_name=SEARCH_BUFFER):
+def load_vi_search_bindings(registry, get_vi_state, get_search_state=None,
+                            filter=None, search_buffer_name=SEARCH_BUFFER):
     assert callable(get_vi_state)  # Callable that takes a CLI and returns a ViState.
+    assert get_search_state is None or callable(get_search_state)
+
+    if not get_search_state:
+        def get_search_state(cli): return cli.search_state
 
     has_focus = filters.HasFocus(search_buffer_name)
     navigation_mode = ~has_focus & (ViStateFilter(get_vi_state, InputMode.NAVIGATION) | filters.HasSelection())
@@ -1273,7 +1283,7 @@ def load_vi_search_bindings(registry, get_vi_state, filter=None, search_buffer_n
         Vi-style forward search.
         """
         # Set the ViState.
-        event.cli.search_state.direction = IncrementalSearchDirection.FORWARD
+        get_search_state(event.cli).direction = IncrementalSearchDirection.FORWARD
         get_vi_state(event.cli).input_mode = InputMode.INSERT
 
         # Focus search buffer.
@@ -1286,7 +1296,7 @@ def load_vi_search_bindings(registry, get_vi_state, filter=None, search_buffer_n
         Vi-style backward search.
         """
         # Set the ViState.
-        event.cli.search_state.direction = IncrementalSearchDirection.BACKWARD
+        get_search_state(event.cli).direction = IncrementalSearchDirection.BACKWARD
 
         # Focus search buffer.
         event.cli.focus_stack.push(search_buffer_name)
@@ -1302,10 +1312,10 @@ def load_vi_search_bindings(registry, get_vi_state, filter=None, search_buffer_n
 
         # Update search state.
         if search_buffer.text:
-            event.cli.search_state.text = search_buffer.text
+            get_search_state(event.cli).text = search_buffer.text
 
         # Apply search.
-        input_buffer.apply_search(event.cli.search_state)
+        input_buffer.apply_search(get_search_state(event.cli))
 
         # Add query to history of search line.
         search_buffer.append_to_history()
diff --git a/prompt_toolkit/key_binding/manager.py b/prompt_toolkit/key_binding/manager.py
index b281d4dc..94f855af 100644
--- a/prompt_toolkit/key_binding/manager.py
+++ b/prompt_toolkit/key_binding/manager.py
@@ -41,6 +41,7 @@ class KeyBindingManager(object):
     :param enable_all: Filter to enable (or disable) all bindings.
     """
     def __init__(self, registry=None, enable_vi_mode=False, get_vi_state=None,
+                 get_search_state=None,
                  enable_abort_and_exit_bindings=False,
                  enable_system_bindings=False, enable_search=False,
                  enable_open_in_editor=False, enable_extra_page_navigation=False,
@@ -49,6 +50,7 @@ class KeyBindingManager(object):
 
         assert registry is None or isinstance(registry, Registry)
         assert get_vi_state is None or callable(get_vi_state)
+        assert get_search_state is None or callable(get_search_state)
 
         # Create registry.
         self.registry = registry or Registry()
@@ -91,7 +93,9 @@ class KeyBindingManager(object):
             self.registry, enable_emacs_mode & enable_open_in_editor & enable_all)
 
         load_emacs_search_bindings(
-            self.registry, enable_emacs_mode & enable_search & enable_all)
+            self.registry,
+            filter=enable_emacs_mode & enable_search & enable_all,
+            get_search_state=get_search_state)
 
         load_emacs_system_bindings(
             self.registry, enable_emacs_mode & enable_system_bindings & enable_all)
@@ -101,8 +105,10 @@ class KeyBindingManager(object):
             enable_emacs_mode & enable_extra_page_navigation & enable_all)
 
         # Load Vi bindings.
-        load_vi_bindings(self.registry, self.get_vi_state, enable_visual_key=~enable_open_in_editor,
-                         filter=enable_vi_mode & enable_all)
+        load_vi_bindings(
+            self.registry, self.get_vi_state, enable_visual_key=~enable_open_in_editor,
+            filter=enable_vi_mode & enable_all,
+            get_search_state=get_search_state)
 
         load_vi_open_in_editor_bindings(
             self.registry, self.get_vi_state,
@@ -110,7 +116,8 @@ class KeyBindingManager(object):
 
         load_vi_search_bindings(
             self.registry, self.get_vi_state,
-            enable_vi_mode & enable_search & enable_all)
+            filter=enable_vi_mode & enable_search & enable_all,
+            get_search_state=get_search_state)
 
         load_vi_system_bindings(
             self.registry, self.get_vi_state,