From f5363d3c2a4ab73f915550a61711b9376f5bbf6b Mon Sep 17 00:00:00 2001
From: James Trew <66286082+jamestrew@users.noreply.github.com>
Date: Sun, 6 Aug 2023 16:46:34 -0400
Subject: [PATCH] feat(diagnostics): add sort_by option (#2632)

* feat(diagnostics): add `sort_by` option

* [docgen] Update doc/telescope.txt
skip-checks: true

---------

Co-authored-by: Github Actions <actions@github>
---
 doc/telescope.txt                       |  6 +++
 lua/telescope/builtin/__diagnostics.lua | 65 +++++++++++++++++--------
 lua/telescope/builtin/init.lua          |  4 ++
 3 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/doc/telescope.txt b/doc/telescope.txt
index 90416ebe14..f45c8f9313 100644
--- a/doc/telescope.txt
+++ b/doc/telescope.txt
@@ -1732,6 +1732,9 @@ builtin.diagnostics({opts})                  *telescope.builtin.diagnostics()*
     - Default keymaps:
       - `<C-l>`: show autocompletion menu to prefilter your query with the
         diagnostic you want to see (i.e. `:warning:`)
+    - sort_by option:
+      - "buffer": order by bufnr (prioritizing current bufnr), severity, lnum
+      - "severity": order by severity, bufnr (prioritizing current bufnr), lnum
 
 
     Parameters: ~
@@ -1763,6 +1766,9 @@ builtin.diagnostics({opts})                  *telescope.builtin.diagnostics()*
                                                 specific namespace
         {disable_coordinates} (boolean)         don't show the line & row
                                                 numbers (default: false)
+        {sort_by}             (string)          sort order of the diagnostics
+                                                results; see above notes
+                                                (default: "buffer")
 
 
 
diff --git a/lua/telescope/builtin/__diagnostics.lua b/lua/telescope/builtin/__diagnostics.lua
index ad50cfa81b..e29be31573 100644
--- a/lua/telescope/builtin/__diagnostics.lua
+++ b/lua/telescope/builtin/__diagnostics.lua
@@ -6,6 +6,50 @@ local utils = require "telescope.utils"
 
 local diagnostics = {}
 
+local sorting_comparator = function(opts)
+  local current_buf = vim.api.nvim_get_current_buf()
+  local comparators = {
+    -- sort results by bufnr (prioritize cur buf), severity, lnum
+    buffer = function(a, b)
+      if a.bufnr == b.bufnr then
+        if a.type == b.type then
+          return a.lnum < b.lnum
+        else
+          return a.type < b.type
+        end
+      else
+        if a.bufnr == current_buf then
+          return true
+        end
+        if b.bufnr == current_buf then
+          return false
+        end
+        return a.bufnr < b.bufnr
+      end
+    end,
+    severity = function(a, b)
+      if a.type < b.type then
+        return true
+      elseif a.type > b.type then
+        return false
+      end
+
+      if a.bufnr == b.bufnr then
+        return a.lnum < b.lnum
+      elseif a.bufnr == current_buf then
+        return true
+      elseif b.bufnr == current_buf then
+        return false
+      else
+        return a.bufnr < b.bufnr
+      end
+    end,
+  }
+
+  local sort_by = vim.F.if_nil(opts.sort_by, "buffer")
+  return comparators[sort_by]
+end
+
 local convert_diagnostic_type = function(severities, severity)
   -- convert from string to int
   if type(severity) == "string" then
@@ -20,7 +64,6 @@ local diagnostics_to_tbl = function(opts)
   opts = vim.F.if_nil(opts, {})
   local items = {}
   local severities = vim.diagnostic.severity
-  local current_buf = vim.api.nvim_get_current_buf()
 
   opts.severity = convert_diagnostic_type(severities, opts.severity)
   opts.severity_limit = convert_diagnostic_type(severities, opts.severity_limit)
@@ -77,25 +120,7 @@ local diagnostics_to_tbl = function(opts)
     end
   end
 
-  -- sort results by bufnr (prioritize cur buf), severity, lnum
-  table.sort(items, function(a, b)
-    if a.bufnr == b.bufnr then
-      if a.type == b.type then
-        return a.lnum < b.lnum
-      else
-        return a.type < b.type
-      end
-    else
-      -- prioritize for current bufnr
-      if a.bufnr == current_buf then
-        return true
-      end
-      if b.bufnr == current_buf then
-        return false
-      end
-      return a.bufnr < b.bufnr
-    end
-  end)
+  table.sort(items, sorting_comparator(opts))
 
   return items
 end
diff --git a/lua/telescope/builtin/init.lua b/lua/telescope/builtin/init.lua
index 99b39710fa..a02b58acb4 100644
--- a/lua/telescope/builtin/init.lua
+++ b/lua/telescope/builtin/init.lua
@@ -521,6 +521,9 @@ builtin.lsp_dynamic_workspace_symbols = require_on_exported_call("telescope.buil
 ---   - `All severity flags can be passed as `string` or `number` as per `:vim.diagnostic.severity:`
 --- - Default keymaps:
 ---   - `<C-l>`: show autocompletion menu to prefilter your query with the diagnostic you want to see (i.e. `:warning:`)
+--- - sort_by option:
+---   - "buffer": order by bufnr (prioritizing current bufnr), severity, lnum
+---   - "severity": order by severity, bufnr (prioritizing current bufnr), lnum
 ---@param opts table: options to pass to the picker
 ---@field bufnr number|nil: Buffer number to get diagnostics from. Use 0 for current buffer or nil for all buffers
 ---@field severity string|number: filter diagnostics by severity name (string) or id (number)
@@ -532,6 +535,7 @@ builtin.lsp_dynamic_workspace_symbols = require_on_exported_call("telescope.buil
 ---@field line_width number: set length of diagnostic entry text in Results
 ---@field namespace number: limit your diagnostics to a specific namespace
 ---@field disable_coordinates boolean: don't show the line & row numbers (default: false)
+---@field sort_by string: sort order of the diagnostics results; see above notes (default: "buffer")
 builtin.diagnostics = require_on_exported_call("telescope.builtin.__diagnostics").get
 
 local apply_config = function(mod)