feat: select_all, toggle_all and drop_all actions (#931)

5 files changed, 219 insertions, 4 deletions

diff --git a/doc/telescope.txt b/doc/telescope.txt
index 37f2b8d..b3f0956 100644
--- a/doc/telescope.txt
+++ b/doc/telescope.txt
@@ -371,6 +371,34 @@ actions.toggle_selection({prompt_bufnr})          *actions.toggle_selection()*
         {prompt_bufnr} (number)  The prompt bufnr
 
 
+actions.select_all({prompt_bufnr})                      *actions.select_all()*
+    Multi select all entries.
+    - Note: selected entries may include results not visible in the results
+      popup.
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.drop_all({prompt_bufnr})                          *actions.drop_all()*
+    Drop all entries from the current multi selection.
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.toggle_all({prompt_bufnr})                      *actions.toggle_all()*
+    Toggle multi selection for all entries.
+    - Note: toggled entries may include results not visible in the results
+      popup.
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
 actions.git_create_branch({prompt_bufnr})        *actions.git_create_branch()*
     Create and checkout a new git branch if it doesn't already exist
 
@@ -429,7 +457,7 @@ actions.git_rebase_branch({prompt_bufnr})        *actions.git_rebase_branch()*
         {prompt_bufnr} (number)  The prompt bufnr
 
 
-actions.git_checkout_current_buffer({prompt_bufnr})*actions.git_checkout_current_buffer()*
+actions.git_staging_toggle({prompt_bufnr})      *actions.git_staging_toggle()*
     Stage/unstage selected file
 
 
@@ -1196,6 +1224,70 @@ builtin.lsp_workspace_diagnostics({opts})*builtin.lsp_workspace_diagnostics()*
 
 
 ================================================================================
+                                                       *telescope.actions.utils*
+
+Utilities to wrap functions around picker selections and entries.
+
+Generally used from within other |telescope.actions|
+
+utils.map_entries({prompt_bufnr}, {f})                   *utils.map_entries()*
+    Apply `f` to the entries of the current picker.
+    - Notes:
+      - Mapped entries may include results not visible in the results popup.
+      - Indices are 1-indexed, whereas rows are 0-indexed.
+    - Warning: `map_entries` has no return value.
+      - The below example showcases how to collect results
+    Usage:
+        local action_state = require "telescope.actions.state"
+        local action_utils = require "telescope.actions.utils"
+        function entry_value_by_row()
+          local prompt_bufnr = vim.api.nvim_get_current_buf()
+          local current_picker = action_state.get_current_picker(prompt_bufnr)
+          local results = {}
+            action_utils.map_entries(prompt_bufnr, function(entry, index, row)
+            results[row] = entry.value
+          end)
+          return results
+        end
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)    The prompt bufnr
+        {f}            (function)  Function to map onto entries of picker that
+                                   takes (entry, index, row) as viable
+                                   arguments
+
+
+utils.map_selections({prompt_bufnr}, {f})             *utils.map_selections()*
+    Apply `f` to the multi selections of the current picker and return a table
+    of mapped selections.
+    - Notes:
+      - Mapped selections may include results not visible in the results popup.
+      - Selected entries are returned in order of their selection.
+    - Warning: `map_selections` has no return value.
+      - The below example showcases how to collect results
+    Usage:
+        local action_state = require "telescope.actions.state"
+        local action_utils = require "telescope.actions.utils"
+        function selection_by_index()
+          local prompt_bufnr = vim.api.nvim_get_current_buf()
+          local current_picker = action_state.get_current_picker(prompt_bufnr)
+          local results = {}
+            action_utils.map_selections(prompt_bufnr, function(entry, index)
+            results[index] = entry.value
+          end)
+          return results
+        end
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)    The prompt bufnr
+        {f}            (function)  Function to map onto selection of picker
+                                   that takes (selection) as a viable argument
+
+
+
+================================================================================
                                                        *telescope.actions.state*
 
 Functions to be used to determine the current state of telescope.
diff --git a/lua/telescope/actions/init.lua b/lua/telescope/actions/init.lua
index d9b3096..af4c1c8 100644
--- a/lua/telescope/actions/init.lua
+++ b/lua/telescope/actions/init.lua
@@ -15,6 +15,7 @@ local utils = require('telescope.utils')
 local p_scroller = require('telescope.pickers.scroller')
 
 local action_state = require('telescope.actions.state')
+local action_utils = require('telescope.actions.utils')
 local action_set = require('telescope.actions.set')
 
 local transform_mod = require('telescope.actions.mt').transform_mod
@@ -143,6 +144,46 @@ function actions.toggle_selection(prompt_bufnr)
   current_picker:toggle_selection(current_picker:get_selection_row())
 end
 
+--- Multi select all entries.
+--- - Note: selected entries may include results not visible in the results popup.
+---@param prompt_bufnr number: The prompt bufnr
+function actions.select_all(prompt_bufnr)
+  local current_picker = action_state.get_current_picker(prompt_bufnr)
+  action_utils.map_entries(prompt_bufnr, function(entry, _, row)
+    if not current_picker._multi:is_selected(entry) then
+      current_picker._multi:add(entry)
+      if current_picker:can_select_row(row) then
+        current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
+      end
+    end
+  end)
+end
+
+--- Drop all entries from the current multi selection.
+---@param prompt_bufnr number: The prompt bufnr
+function actions.drop_all(prompt_bufnr)
+  local current_picker = action_state.get_current_picker(prompt_bufnr)
+  action_utils.map_entries(prompt_bufnr, function(entry, _, row)
+    current_picker._multi:drop(entry)
+    if current_picker:can_select_row(row) then
+      current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
+    end
+  end)
+end
+
+--- Toggle multi selection for all entries.
+--- - Note: toggled entries may include results not visible in the results popup.
+---@param prompt_bufnr number: The prompt bufnr
+function actions.toggle_all(prompt_bufnr)
+  local current_picker = action_state.get_current_picker(prompt_bufnr)
+  action_utils.map_entries(prompt_bufnr, function(entry, _, row)
+    current_picker._multi:toggle(entry)
+    if current_picker:can_select_row(row) then
+      current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
+    end
+  end)
+end
+
 function actions.preview_scrolling_up(prompt_bufnr)
   action_set.scroll_previewer(prompt_bufnr, -1)
 end
@@ -464,8 +505,6 @@ actions.git_rebase_branch = function(prompt_bufnr)
   end
 end
 
---- Stage/unstage selected file
----@param prompt_bufnr number: The prompt bufnr
 actions.git_checkout_current_buffer = function(prompt_bufnr)
   local cwd = actions.get_current_picker(prompt_bufnr).cwd
   local selection = actions.get_selected_entry()
@@ -473,6 +512,8 @@ actions.git_checkout_current_buffer = function(prompt_bufnr)
   utils.get_os_command_output({ 'git', 'checkout', selection.value, '--', selection.file }, cwd)
 end
 
+--- Stage/unstage selected file
+---@param prompt_bufnr number: The prompt bufnr
 actions.git_staging_toggle = function(prompt_bufnr)
   local cwd = action_state.get_current_picker(prompt_bufnr).cwd
   local selection = action_state.get_selected_entry()
diff --git a/lua/telescope/actions/utils.lua b/lua/telescope/actions/utils.lua
new file mode 100644
index 0000000..73045e7
--- /dev/null
+++ b/lua/telescope/actions/utils.lua
@@ -0,0 +1,81 @@
+---@tag telescope.actions.utils
+
+---@brief [[
+--- Utilities to wrap functions around picker selections and entries.
+---
+--- Generally used from within other |telescope.actions|
+---@brief ]]
+
+local action_state = require('telescope.actions.state')
+
+local utils = {}
+
+--- Apply `f` to the entries of the current picker.
+--- - Notes:
+---   - Mapped entries may include results not visible in the results popup.
+---   - Indices are 1-indexed, whereas rows are 0-indexed.
+--- - Warning: `map_entries` has no return value.
+---   - The below example showcases how to collect results
+--- <pre>
+--- Usage:
+---     local action_state = require "telescope.actions.state"
+---     local action_utils = require "telescope.actions.utils"
+---     function entry_value_by_row()
+---       local prompt_bufnr = vim.api.nvim_get_current_buf()
+---       local current_picker = action_state.get_current_picker(prompt_bufnr)
+---       local results = {}
+---         action_utils.map_entries(prompt_bufnr, function(entry, index, row)
+---         results[row] = entry.value
+---       end)
+---       return results
+---     end
+--- </pre>
+---@param prompt_bufnr number: The prompt bufnr
+---@param f function: Function to map onto entries of picker that takes (entry, index, row) as viable arguments
+function utils.map_entries(prompt_bufnr, f)
+  vim.validate{
+    f = {f, "function"}
+  }
+  local current_picker = action_state.get_current_picker(prompt_bufnr)
+  local index = 1
+    -- indices are 1-indexed, rows are 0-indexed
+  for entry in current_picker.manager:iter() do
+    local row = current_picker:get_row(index)
+    f(entry, index, row)
+    index = index + 1
+  end
+end
+
+--- Apply `f` to the multi selections of the current picker and return a table of mapped selections.
+--- - Notes:
+---   - Mapped selections may include results not visible in the results popup.
+---   - Selected entries are returned in order of their selection.
+--- - Warning: `map_selections` has no return value.
+---   - The below example showcases how to collect results
+--- <pre>
+--- Usage:
+---     local action_state = require "telescope.actions.state"
+---     local action_utils = require "telescope.actions.utils"
+---     function selection_by_index()
+---       local prompt_bufnr = vim.api.nvim_get_current_buf()
+---       local current_picker = action_state.get_current_picker(prompt_bufnr)
+---       local results = {}
+---         action_utils.map_selections(prompt_bufnr, function(entry, index)
+---         results[index] = entry.value
+---       end)
+---       return results
+---     end
+--- </pre>
+---@param prompt_bufnr number: The prompt bufnr
+---@param f function: Function to map onto selection of picker that takes (selection) as a viable argument
+function utils.map_selections(prompt_bufnr, f)
+  vim.validate{
+    f = {f, "function"}
+  }
+  local current_picker = action_state.get_current_picker(prompt_bufnr)
+  for _, selection in ipairs(current_picker:get_multi_selection()) do
+    f(selection)
+  end
+end
+
+return utils
diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index e4c73e2..a60d2df 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -264,7 +264,7 @@ function Picker:can_select_row(row)
   if self.sorting_strategy == 'ascending' then
     return row <= self.manager:num_results()
   else
-    return row <= self.max_results and row >= self.max_results - self.manager:num_results()
+    return row >= 0 and row <= self.max_results and row >= self.max_results - self.manager:num_results()
   end
 end
 
diff --git a/scripts/gendocs.lua b/scripts/gendocs.lua
index cd53a30..1533bd5 100644
--- a/scripts/gendocs.lua
+++ b/scripts/gendocs.lua
@@ -15,6 +15,7 @@ docs.test = function()
     "./lua/telescope/actions/init.lua",
     "./lua/telescope/actions/state.lua",
     "./lua/telescope/actions/set.lua",
+    "./lua/telescope/actions/utils.lua",
     "./lua/telescope/previewers/init.lua",
     "./lua/telescope/config/resolve.lua",
     "./lua/telescope/themes.lua",