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

3 files changed, 125 insertions, 3 deletions

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