feat: add filtering by symbol(s) to lsp_{document,workspace}_symbols (#903)

3 files changed, 89 insertions, 6 deletions

diff --git a/lua/telescope/builtin/init.lua b/lua/telescope/builtin/init.lua
index cdf6b71..17cdf88 100644
--- a/lua/telescope/builtin/init.lua
+++ b/lua/telescope/builtin/init.lua
@@ -323,6 +323,7 @@ builtin.lsp_range_code_actions = require('telescope.builtin.lsp').range_code_act
 ---   - `<C-l>`: show autocompletion menu to prefilter your query by type of symbol you want to see (i.e. `:variable:`)
 ---@param opts table: options to pass to the picker
 ---@field ignore_filename type: string with file to ignore
+---@field symbols string|table: filter results by symbol kind(s)
 builtin.lsp_document_symbols = require('telescope.builtin.lsp').document_symbols
 
 --- Lists LSP document symbols in the current workspace
@@ -332,6 +333,7 @@ builtin.lsp_document_symbols = require('telescope.builtin.lsp').document_symbols
 ---@field shorten_path boolean: if true, makes file paths shown in picker use one letter for folders (default is false)
 ---@field ignore_filename string: file(s) to ignore
 ---@field hide_filename boolean: if true, hides the name of the file in the current picker (default is false)
+---@field symbols string|table: filter results by symbol kind(s)
 builtin.lsp_workspace_symbols = require('telescope.builtin.lsp').workspace_symbols
 
 --- Dynamically lists LSP for all workspace symbols
@@ -349,8 +351,8 @@ builtin.lsp_dynamic_workspace_symbols = require('telescope.builtin.lsp').dynamic
 ---@param opts table: options to pass to the picker
 ---@field hide_filename boolean: if true, hides the name of the file in the current picker (default is false)
 ---@field severity string|number: filter diagnostics by severity name (string) or id (number)
----@field severity_limit string|number: filter for diagnostics equal or more severe wrt severity name (string) or id (number)
----@field severity_bound string|number: filter for diagnostics equal or less severe wrt severity name (string) or id (number)
+---@field severity_limit string|number: keep diagnostics equal or more severe wrt severity name (string) or id (number)
+---@field severity_bound string|number: keep diagnostics equal or less severe wrt severity name (string) or id (number)
 ---@field no_sign bool: hide LspDiagnosticSigns from Results (default is false)
 ---@field line_width number: set length of diagnostic entry text in Results
 builtin.lsp_document_diagnostics = require('telescope.builtin.lsp').diagnostics
@@ -363,8 +365,8 @@ builtin.lsp_document_diagnostics = require('telescope.builtin.lsp').diagnostics
 ---@param opts table: options to pass to the picker
 ---@field hide_filename boolean: if true, hides the name of the file in the current picker (default is false)
 ---@field severity string|number: filter diagnostics by severity name (string) or id (number)
----@field severity_limit string|number: filter for diagnostics equal or more severe wrt severity name (string) or id (number)
----@field severity_bound string|number: filter for diagnostics equal or less severe wrt severity name (string) or id (number)
+---@field severity_limit string|number: keep diagnostics equal or more severe wrt severity name (string) or id (number)
+---@field severity_bound string|number: keep diagnostics equal or less severe wrt severity name (string) or id (number)
 ---@field no_sign bool: hide LspDiagnosticSigns from Results (default is false)
 ---@field line_width number: set length of diagnostic entry text in Results
 builtin.lsp_workspace_diagnostics = require('telescope.builtin.lsp').workspace_diagnostics
diff --git a/lua/telescope/builtin/lsp.lua b/lua/telescope/builtin/lsp.lua
index 6a3492b..532f221 100644
--- a/lua/telescope/builtin/lsp.lua
+++ b/lua/telescope/builtin/lsp.lua
@@ -107,6 +107,12 @@ lsp.document_symbols = function(opts)
     vim.list_extend(locations, vim.lsp.util.symbols_to_items(server_results.result, 0) or {})
   end
 
+  locations = utils.filter_symbols(locations, opts)
+  if locations == nil then
+    -- error message already printed in `utils.filter_symbols`
+    return
+  end
+
   if vim.tbl_isempty(locations) then
     return
   end
@@ -263,6 +269,12 @@ lsp.workspace_symbols = function(opts)
     end
   end
 
+  locations = utils.filter_symbols(locations, opts)
+  if locations == nil then
+    -- error message already printed in `utils.filter_symbols`
+    return
+  end
+
   if vim.tbl_isempty(locations) then
     print("No results from workspace/symbol. Maybe try a different query: " ..
       "Telescope lsp_workspace_symbols query=example")
diff --git a/lua/telescope/utils.lua b/lua/telescope/utils.lua
index 727330d..d9e4ec0 100644
--- a/lua/telescope/utils.lua
+++ b/lua/telescope/utils.lua
@@ -82,6 +82,75 @@ utils.quickfix_items_to_entries = function(locations)
   return results
 end
 
+utils.filter_symbols = function(results, opts)
+  if opts.symbols == nil then
+    return results
+  end
+  local valid_symbols = vim.tbl_map(string.lower, vim.lsp.protocol.SymbolKind)
+
+  local filtered_symbols = {}
+  if type(opts.symbols) == "string" then
+    opts.symbols = string.lower(opts.symbols)
+    if vim.tbl_contains(valid_symbols, opts.symbols) then
+      for _, result in ipairs(results) do
+        if string.lower(result.kind) == opts.symbols then
+          table.insert(filtered_symbols, result)
+        end
+      end
+    else
+      print(string.format("%s is not a valid symbol per `vim.lsp.protocol.SymbolKind`", opts.symbols))
+    end
+  elseif type(opts.symbols) == "table" then
+    opts.symbols = vim.tbl_map(string.lower, opts.symbols)
+    local mismatched_symbols = {}
+    for _, symbol in ipairs(opts.symbols) do
+      if vim.tbl_contains(valid_symbols, symbol) then
+        for _, result in ipairs(results) do
+          if string.lower(result.kind) == symbol then
+            table.insert(filtered_symbols, result)
+          end
+        end
+      else
+        table.insert(mismatched_symbols, symbol)
+        mismatched_symbols = table.concat(mismatched_symbols, ", ")
+        print(string.format("%s are not valid symbols per `vim.lsp.protocol.SymbolKind`", mismatched_symbols))
+      end
+    end
+  else
+    print("Please pass filtering symbols as either a string or a list of strings")
+    return
+  end
+
+  local current_buf = vim.api.nvim_get_current_buf()
+  if not vim.tbl_isempty(filtered_symbols) then
+    -- filter adequately for workspace symbols
+    local filename_to_bufnr = {}
+    for _, symbol in ipairs(filtered_symbols) do
+      if filename_to_bufnr[symbol.filename] == nil then
+        filename_to_bufnr[symbol.filename] = vim.uri_to_bufnr(vim.uri_from_fname(symbol.filename))
+      end
+      symbol['bufnr'] = filename_to_bufnr[symbol.filename]
+    end
+    table.sort(filtered_symbols, function(a, b)
+      if a.bufnr == b.bufnr then
+        return a.lnum < b.lnum
+      end
+      if a.bufnr == current_buf then
+        return true
+      end
+      if b.bufnr == current_buf then
+        return false
+      end
+      return a.bufnr < b.bufnr
+    end)
+  return filtered_symbols
+  end
+  -- only account for string|table as function otherwise already printed message and returned nil
+  local symbols = type(opts.symbols) == 'string' and opts.symbols or table.concat(opts.symbols, ', ')
+  print(string.format("%s symbol(s) were not part of the query results", symbols))
+  return
+end
+
 local convert_diagnostic_type = function(severity)
   -- convert from string to int
   if type(severity) == 'string' then
@@ -115,7 +184,7 @@ utils.diagnostics_to_tbl = function(opts)
   opts.severity_bound = convert_diagnostic_type(opts.severity_bound)
 
   local validate_severity = 0
-  for _, v in pairs({opts.severity, opts.severity_limit, opts.severity_bound}) do
+  for _, v in ipairs({opts.severity, opts.severity_limit, opts.severity_bound}) do
     if v ~= nil then
       validate_severity = validate_severity + 1
     end
@@ -149,7 +218,7 @@ utils.diagnostics_to_tbl = function(opts)
   local buffer_diags = opts.get_all and vim.lsp.diagnostic.get_all() or
     {[current_buf] = vim.lsp.diagnostic.get(current_buf, opts.client_id)}
   for bufnr, diags in pairs(buffer_diags) do
-    for _, diag in pairs(diags) do
+    for _, diag in ipairs(diags) do
       -- workspace diagnostics may include empty tables for unused bufnr
       if not vim.tbl_isempty(diag) then
         if filter_diag_severity(opts, diag.severity) then