feat(diagnostic): add 'prefix' option to open_float (#16321)

The 'prefix' option accepts a function or a string that is used to add a
prefix string to each diagnostic displayed in the floating window.

3 files changed, 113 insertions, 5 deletions

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index e2469b4f4e..d53aae510e 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -384,6 +384,17 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
                                      a diagnostic as input and returns a
                                      string. The return value is the text used
                                      to display the diagnostic.
+                                   • prefix: (function or string) Prefix each
+                                     diagnostic in the floating window. If a
+                                     function, it must have the signature
+                                     (diagnostic, i, total) -> string, where
+                                     {i} is the index of the diagnostic being
+                                     evaluated and {total} is the total number
+                                     of diagnostics displayed in the window.
+                                     The returned string is prepended to each
+                                     diagnostic in the window. Otherwise, if
+                                     {prefix} is a string, it is prepended to
+                                     each diagnostic.
 
                                  • update_in_insert: (default false) Update
                                    diagnostics in Insert mode (if false,
@@ -614,6 +625,9 @@ open_float({bufnr}, {opts})                      *vim.diagnostic.open_float()*
                                return value is the text used to display the
                                diagnostic. Overrides the setting from
                                |vim.diagnostic.config()|.
+                             • prefix: (function or string) Prefix each
+                               diagnostic in the floating window. Overrides
+                               the setting from |vim.diagnostic.config()|.
 
                 Return: ~
                     tuple ({float_bufnr}, {win_id})
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index d0f729f3b9..b30a678eeb 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -1,3 +1,5 @@
+local if_nil = vim.F.if_nil
+
 local M = {}
 
 M.severity = {
@@ -65,7 +67,7 @@ end
 local function prefix_source(source, diagnostics)
   vim.validate { source = {source, function(v)
     return v == "always" or v == "if_many"
-  end, "Invalid value for option 'source'" } }
+  end, "'always' or 'if_many'" } }
 
   if source == "if_many" then
     local sources = {}
@@ -537,6 +539,12 @@ end
 ---                            the message. One of "always" or "if_many".
 ---                  * format: (function) A function that takes a diagnostic as input and returns a
 ---                            string. The return value is the text used to display the diagnostic.
+---                  * prefix: (function or string) Prefix each diagnostic in the floating window. If
+---                            a function, it must have the signature (diagnostic, i, total) -> string,
+---                            where {i} is the index of the diagnostic being evaluated and {total} is
+---                            the total number of diagnostics displayed in the window. The returned
+---                            string is prepended to each diagnostic in the window. Otherwise,
+---                            if {prefix} is a string, it is prepended to each diagnostic.
 ---       - update_in_insert: (default false) Update diagnostics in Insert mode (if false,
 ---                           diagnostics are updated on InsertLeave)
 ---       - severity_sort: (default false) Sort diagnostics by severity. This affects the order in
@@ -1140,6 +1148,8 @@ end
 ---            - format: (function) A function that takes a diagnostic as input and returns a
 ---                      string. The return value is the text used to display the diagnostic.
 ---                      Overrides the setting from |vim.diagnostic.config()|.
+---            - prefix: (function or string) Prefix each diagnostic in the floating window.
+---                      Overrides the setting from |vim.diagnostic.config()|.
 ---@return tuple ({float_bufnr}, {win_id})
 function M.open_float(bufnr, opts)
   vim.validate {
@@ -1224,8 +1234,17 @@ function M.open_float(bufnr, opts)
     diagnostics = prefix_source(opts.source, diagnostics)
   end
 
+  local prefix_opt = if_nil(opts.prefix, (scope == "cursor" and #diagnostics <= 1) and "" or function(_, i)
+      return string.format("%d. ", i)
+  end)
+  if prefix_opt then
+    vim.validate { prefix = { prefix_opt, function(v)
+      return type(v) == "string" or type(v) == "function"
+    end, "'string' or 'function'" } }
+  end
+
   for i, diagnostic in ipairs(diagnostics) do
-    local prefix = string.format("%d. ", i)
+    local prefix = type(prefix_opt) == "string" and prefix_opt or prefix_opt(diagnostic, i, #diagnostics)
     local hiname = floating_highlight_map[diagnostic.severity]
     local message_lines = vim.split(diagnostic.message, '\n')
     table.insert(lines, prefix..message_lines[1])
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index bb2d39c8b2..ec53916024 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -1081,7 +1081,7 @@ describe('vim.diagnostic', function()
 
     it('can show diagnostics from a specific position', function()
       -- Using cursor position
-      eq({'1. Syntax error'}, exec_lua [[
+      eq({'Syntax error'}, exec_lua [[
         local diagnostics = {
           make_error("Syntax error", 1, 1, 1, 2),
           make_warning("Some warning", 1, 3, 1, 4),
@@ -1096,7 +1096,7 @@ describe('vim.diagnostic', function()
       ]])
 
       -- With specified position
-      eq({'1. Some warning'}, exec_lua [[
+      eq({'Some warning'}, exec_lua [[
         local diagnostics = {
           make_error("Syntax error", 1, 1, 1, 2),
           make_warning("Some warning", 1, 3, 1, 4),
@@ -1111,7 +1111,7 @@ describe('vim.diagnostic', function()
       ]])
 
       -- With column position past the end of the line. #16062
-      eq({'1. Syntax error'}, exec_lua [[
+      eq({'Syntax error'}, exec_lua [[
         local first_line_len = #vim.api.nvim_buf_get_lines(diagnostic_bufnr, 0, 1, true)[1]
         local diagnostics = {
           make_error("Syntax error", 0, first_line_len + 1, 1, 0),
@@ -1325,6 +1325,81 @@ describe('vim.diagnostic', function()
       eq(4, count_diagnostics_with_severity("HINT"))
       eq(0, count_diagnostics_with_severity("HINT", "HINT"))
     end)
+
+    it('can add a prefix to diagnostics', function()
+      -- Default is to add a number
+      eq({'1. Syntax error', '2. Some warning'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3),
+          make_warning("Some warning", 1, 1, 1, 3),
+        }
+        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local float_bufnr, winnr = vim.diagnostic.open_float(0, {show_header = false, number = "always"})
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq({'Syntax error', 'Some warning'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3),
+          make_warning("Some warning", 1, 1, 1, 3),
+        }
+        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local float_bufnr, winnr = vim.diagnostic.open_float(0, {show_header = false, prefix = ""})
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq({'1. Syntax error', '2. Some warning'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3),
+          make_warning("Some warning", 1, 1, 1, 3),
+        }
+        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local float_bufnr, winnr = vim.diagnostic.open_float(0, {
+          show_header = false,
+          prefix = function(_, i, total)
+            -- Only show a number if there is more than one diagnostic
+            if total > 1 then
+              return string.format("%d. ", i)
+            end
+            return ""
+          end,
+        })
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq({'Syntax error'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3),
+        }
+        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local float_bufnr, winnr = vim.diagnostic.open_float(0, {
+          show_header = false,
+          prefix = function(_, i, total)
+            -- Only show a number if there is more than one diagnostic
+            if total > 1 then
+              return string.format("%d. ", i)
+            end
+            return ""
+          end,
+        })
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq("Error executing lua: .../diagnostic.lua:0: prefix: expected 'string' or 'function', got 42",
+        pcall_err(exec_lua, [[ vim.diagnostic.open_float(0, { prefix = 42 }) ]]))
+    end)
   end)
 
   describe('setloclist()', function()