feat(pickers): implement group-by.

lua/telescope/config.lua

@@ -1,6 +1,7 @@
 local strings = require "plenary.strings"
 local deprecated = require "telescope.deprecated"
 local sorters = require "telescope.sorters"
+local utils = require "telescope.utils"
 local os_sep = require("plenary.path").path.sep
 local has_win = vim.fn.has "win32" == 1
 
@@ -864,6 +865,88 @@ append(
     Default: require("telescope.previewers").buffer_previewer_maker]]
 )
 
+append(
+  "group_by_options",
+  {
+    defaults = {
+      group_getter = function(field)
+        return function(object)
+          return object[field]
+        end
+      end,
+      header_renderer = function(opts, value)
+        return { { value .. ":", opts.group_by.header_highlights } }
+      end,
+      header_highlights = { "@constant", "Bold" },
+    },
+    fields = {
+      ["filename"] = {
+        header_renderer = function(opts, value)
+          local devicon, devicon_hl = utils.get_devicons(value, opts.group_by.disable_devicons)
+          return {
+            { devicon .. (opts.group_by.disable_devicons and "" or " "), devicon_hl },
+            { utils.transform_path(opts, value) .. ":", opts.group_by.header_highlights },
+          }
+        end,
+      },
+    },
+  },
+  [[
+    This field defines the default behaviors of the group_by picker argument.
+
+    Fields:
+      - defaults:   specifies the default values for fields of the options.
+      - fields:     a dictionary whose keys are field names, and whose values
+                    are the default values for fields of the options when
+                    grouping by the key. Overrides the values in `defaults`
+                    for those keys.
+
+    Fields of the final options object (an object with these fields may also
+    be passed as the `group_by` argument to a picker call):
+      - field:               name of the field to group by.
+      - group_getter:        a function that, once provided with a field name,
+                             returns a function that extracts the group by key
+                             from an object.
+      - header_renderer:     a function that renders the headers for the groups
+                             being displayed.
+      - header_highlights:   highlights to be used for the group_by headers.
+                             Default: `{ "@constant", "Bold" }`.
+      - disable_devicons:    If the group by key is `filename`, controls whether
+                             the group by headers include devicons.
+                             Default: false.]]
+)
+
+-- @param argument string|table: if string, the name of the field to group by.
+--      If a table, contains the overrides for default options.
+function config.resolve_group_by_opts(argument)
+  if argument == nil then
+    return nil
+  end
+
+  local options = argument
+  if type(argument) == "string" then
+    options = config.values.group_by_options.fields[argument] or {}
+    options.field = argument
+  end
+
+  if argument.field ~= nil then
+    options = vim.tbl_deep_extend("keep", options, config.values.group_by_options.fields[argument.field])
+  end
+
+  if config.values.group_by_options ~= nil then
+    options = vim.tbl_deep_extend("keep", options, vim.F.if_nil(config.values.group_by_options.defaults, {}))
+  end
+
+  if options.field ~= nil then
+    options = vim.tbl_deep_extend("keep", options, telescope_defaults.group_by_options[1].fields[options.field])
+  end
+  options = vim.tbl_deep_extend("keep", options, telescope_defaults.group_by_options[1].defaults)
+
+  options._resolved_group_getter = options.group_getter(options.field)
+
+  return options
+end
+
 -- @param user_defaults table: a table where keys are the names of options,
 --    and values are the ones the user wants
 -- @param tele_defaults table: (optional) a table containing all of the defaults

lua/telescope/pickers.lua

@@ -31,6 +31,7 @@ local strdisplaywidth = require("plenary.strings").strdisplaywidth
 local ns_telescope_matching = a.nvim_create_namespace "telescope_matching"
 local ns_telescope_prompt = a.nvim_create_namespace "telescope_prompt"
 local ns_telescope_prompt_prefix = a.nvim_create_namespace "telescope_prompt_prefix"
+local ns_telescope_group_by_headers = a.nvim_create_namespace "telescope_group_by_headers"
 
 ---@class telescope_popup_options
 ---@field border table<1|2|3|4, integer>
@@ -317,6 +318,12 @@ function Picker:new(opts)
     cache_picker = config.resolve_table_opts(opts.cache_picker, vim.deepcopy(config.values.cache_picker)),
 
     __scrolling_limit = tonumber(vim.F.if_nil(opts.temp__scrolling_limit, 250)),
+
+    group_by = config.resolve_group_by_opts(opts.group_by),
+    -- This is a flag testing for whether group_by needs a workaround for neovim/neovim#16166.
+    -- When it does, we insert an additional buffer line at the top of the results buffer, and
+    -- use it to fake a virt_line at the top of the buffer.
+    _group_by_workaround = opts.group_by ~= nil,
   }, self)
 
   obj.create_layout = opts.create_layout or config.values.create_layout or default_create_layout
@@ -362,15 +369,19 @@ function Picker:new(opts)
   return obj
 end
 
+function Picker:_group_by_offset()
+  return self._group_by_workaround and 1 or 0
+end
+
 --- Take an index and get a row.
 ---@note: Rows are 0-indexed, and `index` is 1 indexed (table index)
 ---@param index number: the index in line_manager
 ---@return number: the row for the picker to display in
 function Picker:get_row(index)
   if self.sorting_strategy == "ascending" then
-    return index - 1
+    return index - 1 + self:_group_by_offset()
   else
-    return self.max_results - index
+    return self.max_results + self:_group_by_offset() - index
   end
 end
 
@@ -380,19 +391,19 @@ end
 ---@return number: The index in line_manager
 function Picker:get_index(row)
   if self.sorting_strategy == "ascending" then
-    return row + 1
+    return row + 1 - self:_group_by_offset()
   else
-    return self.max_results - row
+    return self.max_results + self:_group_by_offset() - row
   end
 end
 
 --- Get the row number of the "best" entry
 ---@return number: the number of the "reset" row
 function Picker:get_reset_row()
   if self.sorting_strategy == "ascending" then
-    return 0
+    return 0 + self:_group_by_offset()
   else
-    return self.max_results - 1
+    return self.max_results - 1 + self:_group_by_offset()
   end
 end
 
@@ -418,24 +429,41 @@ function Picker:clear_extra_rows(results_bufnr)
     return
   end
 
-  local worst_line, ok, msg
-  if self.sorting_strategy == "ascending" then
-    local num_results = self.manager:num_results()
-    worst_line = self.max_results - num_results
+  local num_results = self.manager:num_results()
+  local worst_line = self:get_row(num_results)
+  if worst_line <= 0 and num_results ~= 0 and num_results ~= self.max_results then
+    return
+  end
 
-    if worst_line <= 0 then
-      return
+  local ok, msg
+  if self.sorting_strategy == "ascending" then
+    ok, msg = pcall(
+      vim.api.nvim_buf_set_lines,
+      results_bufnr,
+      vim.fn.min { num_results, self.max_results } + self:_group_by_offset(),
+      -1,
+      false,
+      {}
+    )
+    if self.group_by ~= nil then
+      local begin = num_results == 0 and 0 or num_results + self:_group_by_offset()
+      pcall(vim.api.nvim_buf_clear_namespace, results_bufnr, ns_telescope_group_by_headers, begin, -1)
     end
-
-    ok, msg = pcall(vim.api.nvim_buf_set_lines, results_bufnr, num_results, -1, false, {})
   else
-    worst_line = self:get_row(self.manager:num_results())
-    if worst_line <= 0 then
-      return
-    end
-
     local empty_lines = utils.repeated_table(worst_line, "")
     ok, msg = pcall(vim.api.nvim_buf_set_lines, results_bufnr, 0, worst_line, false, empty_lines)
+    if self.group_by ~= nil then
+      pcall(
+        vim.api.nvim_buf_clear_namespace,
+        results_bufnr,
+        ns_telescope_group_by_headers,
+        0,
+        worst_line + self:_group_by_offset()
+      )
+    end
+    if self.group_by ~= nil then
+      self:_set_headers(self.manager:num_results())
+    end
   end
 
   if not ok then
@@ -486,10 +514,14 @@ end
 ---@param row number: the number of the chosen row in the results buffer
 ---@return boolean
 function Picker:can_select_row(row)
+  local war_offset = self:_group_by_offset()
   if self.sorting_strategy == "ascending" then
-    return row <= self.manager:num_results() and row < self.max_results
+    return row <= vim.fn.min { self.manager:num_results(), self.max_results } + war_offset
+      and not (self._group_by_workaround and row == 0)
   else
-    return row >= 0 and row <= self.max_results and row >= self.max_results - self.manager:num_results()
+    return row <= self.max_results + war_offset
+      and row >= war_offset
+      and row >= self.max_results + war_offset - self.manager:num_results()
   end
 end
 
@@ -521,6 +553,9 @@ end
 function Picker:find()
   self:close_existing_pickers()
   self:reset_selection()
+  if self.results_bufnr and a.nvim_buf_is_valid(self.results_bufnr) then
+    a.nvim_buf_clear_namespace(self.results_bufnr, ns_telescope_group_by_headers, 0, -1)
+  end
 
   self.original_win_id = a.nvim_get_current_win()
 
@@ -564,7 +599,13 @@ function Picker:find()
   -- This just lets us stop doing stuff after tons of  things.
   self.max_results = self.__scrolling_limit
 
-  vim.api.nvim_buf_set_lines(self.results_bufnr, 0, self.max_results, false, utils.repeated_table(self.max_results, ""))
+  vim.api.nvim_buf_set_lines(
+    self.results_bufnr,
+    0,
+    self.max_results + self:_group_by_offset(),
+    false,
+    utils.repeated_table(self.max_results + self:_group_by_offset(), "")
+  )
 
   local status_updater = self:get_status_updater(self.prompt_win, self.prompt_bufnr)
   local debounced_status = debounce.throttle_leading(status_updater, 50)
@@ -1017,7 +1058,9 @@ function Picker:set_selection(row)
     return
   end
 
+  row = row - self:_group_by_offset()
   row = self.scroller(self.max_results, self.manager:num_results(), row)
+  row = row + self:_group_by_offset()
 
   if not self:can_select_row(row) then
     -- If the current selected row exceeds number of currently displayed
@@ -1111,11 +1154,24 @@ function Picker:set_selection(row)
 
   self:refresh_previewer()
 
+  if
+    self._group_by_workaround and row == 1
+    or (self.sorting_strategy == "descending" and row + self.manager:num_results() == self.max_results + 1)
+  then
+    vim.api.nvim_win_set_cursor(self.results_win, { row, 0 })
+  end
   vim.api.nvim_win_set_cursor(self.results_win, { row + 1, 0 })
 end
 
 --- Update prefix for entry on a given row
 function Picker:update_prefix(entry, row)
+  if
+    self._group_by_workaround
+    and (row == 0 or (self.sorting_strategy == "descending" and row + self.manager:num_results() == self.max_results))
+  then
+    return
+  end
+
   local prefix = function(sel, multi)
     local t
     if sel then
@@ -1197,6 +1253,43 @@ function Picker:cycle_previewers(next)
   end
 end
 
+function Picker:_set_headers(index, entry)
+  if index == 0 then
+    return
+  end
+  entry = entry or self.manager:get_entry(index)
+
+  local direction = self.sorting_strategy == "ascending" and 1 or -1
+  local row = self:get_row(index)
+
+  local above = self.manager:get_entry(index - direction)
+  if above == nil then
+    a.nvim_buf_clear_namespace(self.results_bufnr, ns_telescope_group_by_headers, row - 1, row + 1)
+    a.nvim_buf_set_extmark(self.results_bufnr, ns_telescope_group_by_headers, row - 1, 0, {
+      virt_text = self.group_by.header_renderer(self, self.group_by._resolved_group_getter(entry)),
+      virt_text_pos = "overlay",
+      virt_text_win_col = 0,
+    })
+  elseif self.group_by._resolved_group_getter(entry) ~= self.group_by._resolved_group_getter(above) then
+    a.nvim_buf_set_extmark(self.results_bufnr, ns_telescope_group_by_headers, row, 0, {
+      virt_lines = { self.group_by.header_renderer(self, self.group_by._resolved_group_getter(entry)) },
+      virt_lines_above = true,
+    })
+  end
+
+  local below = self.manager:get_entry(index + direction)
+  if below ~= nil then
+    local same_group_below = self.group_by._resolved_group_getter(entry) == self.group_by._resolved_group_getter(below)
+    a.nvim_buf_clear_namespace(self.results_bufnr, ns_telescope_group_by_headers, row + 1, row + 2)
+    if not same_group_below then
+      a.nvim_buf_set_extmark(self.results_bufnr, ns_telescope_group_by_headers, row + 1, 0, {
+        virt_lines = { self.group_by.header_renderer(self, self.group_by._resolved_group_getter(below)) },
+        virt_lines_above = true,
+      })
+    end
+  end
+end
+
 --- Handler for when entries are added by `self.manager`
 ---@param index number: the index to add the entry at
 ---@param entry table: the entry that has been added to the manager
@@ -1247,6 +1340,10 @@ function Picker:entry_adder(index, entry, _, insert)
     end
   end
 
+  if not insert and self.group_by ~= nil then
+    vim.api.nvim_buf_clear_namespace(self.results_bufnr, ns_telescope_group_by_headers, row, row + 1)
+  end
+
   local set_ok, msg = pcall(vim.api.nvim_buf_set_lines, self.results_bufnr, row, row + offset, false, { display })
   if set_ok then
     if display_highlights then
@@ -1266,6 +1363,10 @@ function Picker:entry_adder(index, entry, _, insert)
     display = display:gsub("\n", " | ")
     vim.api.nvim_buf_set_lines(self.results_bufnr, row, row + 1, false, { display })
   end
+
+  if self.group_by ~= nil then
+    self:_set_headers(index, entry)
+  end
 end
 
 --- Reset tracked information for this picker
@@ -1418,8 +1519,11 @@ function Picker:get_result_completor(results_bufnr, find_id, prompt, status_upda
 
     if self.sorting_strategy == "descending" then
       local visible_result_rows = vim.api.nvim_win_get_height(self.results_win)
-      vim.api.nvim_win_set_cursor(self.results_win, { self.max_results - visible_result_rows, 1 })
-      vim.api.nvim_win_set_cursor(self.results_win, { self.max_results, 1 })
+      vim.api.nvim_win_set_cursor(
+        self.results_win,
+        { self.max_results + self:_group_by_offset() - visible_result_rows, 1 }
+      )
+      vim.api.nvim_win_set_cursor(self.results_win, { self.max_results + self:_group_by_offset(), 1 })
     end
     self:_on_complete()
   end)
@@ -1574,6 +1678,15 @@ function pickers.on_close_prompt(prompt_bufnr)
     picker.finder:close()
   end
 
+  if picker.group_by ~= nil then
+    a.nvim_buf_clear_namespace(picker.results_bufnr, ns_telescope_group_by_headers, 0, -1)
+    if picker.group_by_workaround_applied then
+      picker.group_by_workaround_applied = false
+      a.nvim_buf_set_lines(picker.results_bufnr, 0, 1, false, {})
+    end
+    picker.group_by = nil
+  end
+
   -- so we dont call close_windows multiple times we clear that autocmd
   vim.api.nvim_clear_autocmds {
     group = "PickerInsert",