Allow searching in Vi selection mode.

2 files changed, 16 insertions, 4 deletions

diff --git a/prompt_toolkit/buffer.py b/prompt_toolkit/buffer.py
index 533309dc..7636b7b6 100644
--- a/prompt_toolkit/buffer.py
+++ b/prompt_toolkit/buffer.py
@@ -1084,7 +1084,10 @@ class Buffer(object):
     def document_for_search(self, search_state):
         """
         Return a :class:`~prompt_toolkit.document.Document` instance that has
-        the text/cursor position for this search, if we would apply it.
+        the text/cursor position for this search, if we would apply it. This
+        will be used in the
+        :class:`~prompt_toolkit.layout.controls.BufferControl` to display
+        feedback while searching.
         """
         search_result = self._search(search_state, include_current_position=True)
 
@@ -1092,7 +1095,15 @@ class Buffer(object):
             return self.document
         else:
             working_index, cursor_position = search_result
-            return Document(self._working_lines[working_index], cursor_position)
+
+            # Keep selection, when `working_index` was not changed.
+            if working_index == self.working_index:
+                selection = self.selection_state
+            else:
+                selection = None
+
+            return Document(self._working_lines[working_index],
+                            cursor_position, selection=selection)
 
     def get_search_position(self, search_state, include_current_position=True, count=1):
         """
diff --git a/prompt_toolkit/key_binding/bindings/vi.py b/prompt_toolkit/key_binding/bindings/vi.py
index d03f6a75..318e8b36 100644
--- a/prompt_toolkit/key_binding/bindings/vi.py
+++ b/prompt_toolkit/key_binding/bindings/vi.py
@@ -1741,9 +1741,10 @@ def load_vi_search_bindings(registry, get_search_state=None,
 
     has_focus = filters.HasFocus(search_buffer_name)
     navigation_mode = ViNavigationMode()
+    selection_mode = ViSelectionMode()
     handle = create_handle_decorator(registry, filter & ViMode())
 
-    @handle('/', filter=navigation_mode)
+    @handle('/', filter=navigation_mode|selection_mode)
     @handle(Keys.ControlS, filter=~has_focus)
     def _(event):
         """
@@ -1756,7 +1757,7 @@ def load_vi_search_bindings(registry, get_search_state=None,
         # Focus search buffer.
         event.cli.push_focus(search_buffer_name)
 
-    @handle('?', filter=navigation_mode)
+    @handle('?', filter=navigation_mode|selection_mode)
     @handle(Keys.ControlR, filter=~has_focus)
     def _(event):
         """