feat(diagnostic): add `suffix` option to `open_float()` (#21130)

Closes #18687

This introduces a `suffix` option to `vim.diagnostic.open_float()` (and
consequently `vim.diagnostic.config()`) that appends some text to each
diagnostic in the float.

It accepts the same types as `prefix`. For multiline diagnostics, the suffix is
only appended to the last line. By default, the suffix will render the
diagnostic error code, if any.

3 files changed, 111 insertions, 21 deletions

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index d30d6c28a0..53046b399f 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -618,6 +618,9 @@ open_float({opts}, {...})                        *vim.diagnostic.open_float()*
                   {prefix} is a string, it is prepended to each diagnostic in
                   the window with no highlight. Overrides the setting from
                   |vim.diagnostic.config()|.
+                • suffix: Same as {prefix}, but appends the text to the
+                  diagnostic instead of prepending it. 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 09eeb617f5..008b00a9a0 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -1220,6 +1220,8 @@ end
 ---                      string, it is prepended to each diagnostic in the window with no
 ---                      highlight.
 ---                      Overrides the setting from |vim.diagnostic.config()|.
+---            - suffix: Same as {prefix}, but appends the text to the diagnostic instead of
+---                      prepending it. Overrides the setting from |vim.diagnostic.config()|.
 ---@return tuple ({float_bufnr}, {win_id})
 function M.open_float(opts, ...)
   -- Support old (bufnr, opts) signature
@@ -1313,11 +1315,11 @@ function M.open_float(opts, ...)
       -- Don't insert any lines for an empty string
       if string.len(if_nil(header[1], '')) > 0 then
         table.insert(lines, header[1])
-        table.insert(highlights, { 0, header[2] or 'Bold' })
+        table.insert(highlights, { hlname = header[2] or 'Bold' })
       end
     elseif #header > 0 then
       table.insert(lines, header)
-      table.insert(highlights, { 0, 'Bold' })
+      table.insert(highlights, { hlname = 'Bold' })
     end
   end
 
@@ -1350,18 +1352,52 @@ function M.open_float(opts, ...)
     end
   end
 
+  local suffix_opt = if_nil(opts.suffix, function(diagnostic)
+    return diagnostic.code and string.format(' [%s]', diagnostic.code) or ''
+  end)
+
+  local suffix, suffix_hl_group
+  if suffix_opt then
+    vim.validate({
+      suffix = {
+        suffix_opt,
+        { 'string', 'table', 'function' },
+        "'string' or 'table' or 'function'",
+      },
+    })
+    if type(suffix_opt) == 'string' then
+      suffix, suffix_hl_group = suffix_opt, 'NormalFloat'
+    elseif type(suffix_opt) == 'table' then
+      suffix, suffix_hl_group = suffix_opt[1] or '', suffix_opt[2] or 'NormalFloat'
+    end
+  end
+
   for i, diagnostic in ipairs(diagnostics) do
     if prefix_opt and type(prefix_opt) == 'function' then
       prefix, prefix_hl_group = prefix_opt(diagnostic, i, #diagnostics)
       prefix, prefix_hl_group = prefix or '', prefix_hl_group or 'NormalFloat'
     end
+    if suffix_opt and type(suffix_opt) == 'function' then
+      suffix, suffix_hl_group = suffix_opt(diagnostic, i, #diagnostics)
+      suffix, suffix_hl_group = suffix or '', suffix_hl_group or 'NormalFloat'
+    end
     local hiname = floating_highlight_map[diagnostic.severity]
     local message_lines = vim.split(diagnostic.message, '\n')
-    table.insert(lines, prefix .. message_lines[1])
-    table.insert(highlights, { #prefix, hiname, prefix_hl_group })
-    for j = 2, #message_lines do
-      table.insert(lines, string.rep(' ', #prefix) .. message_lines[j])
-      table.insert(highlights, { 0, hiname })
+    for j = 1, #message_lines do
+      local pre = j == 1 and prefix or string.rep(' ', #prefix)
+      local suf = j == #message_lines and suffix or ''
+      table.insert(lines, pre .. message_lines[j] .. suf)
+      table.insert(highlights, {
+        hlname = hiname,
+        prefix = {
+          length = j == 1 and #prefix or 0,
+          hlname = prefix_hl_group,
+        },
+        suffix = {
+          length = j == #message_lines and #suffix or 0,
+          hlname = suffix_hl_group,
+        },
+      })
     end
   end
 
@@ -1370,12 +1406,17 @@ function M.open_float(opts, ...)
     opts.focus_id = scope
   end
   local float_bufnr, winnr = require('vim.lsp.util').open_floating_preview(lines, 'plaintext', opts)
-  for i, hi in ipairs(highlights) do
-    local prefixlen, hiname, prefix_hiname = unpack(hi)
-    if prefix_hiname then
-      api.nvim_buf_add_highlight(float_bufnr, -1, prefix_hiname, i - 1, 0, prefixlen)
+  for i, hl in ipairs(highlights) do
+    local line = lines[i]
+    local prefix_len = hl.prefix and hl.prefix.length or 0
+    local suffix_len = hl.suffix and hl.suffix.length or 0
+    if prefix_len > 0 then
+      api.nvim_buf_add_highlight(float_bufnr, -1, hl.prefix.hlname, i - 1, 0, prefix_len)
+    end
+    api.nvim_buf_add_highlight(float_bufnr, -1, hl.hlname, i - 1, prefix_len, #line - suffix_len)
+    if suffix_len > 0 then
+      api.nvim_buf_add_highlight(float_bufnr, -1, hl.suffix.hlname, i - 1, #line - suffix_len, -1)
     end
-    api.nvim_buf_add_highlight(float_bufnr, -1, hiname, i - 1, prefixlen, -1)
   end
 
   return float_bufnr, winnr
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 116a541358..816b359725 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -16,7 +16,7 @@ describe('vim.diagnostic', function()
     exec_lua [[
       require('vim.diagnostic')
 
-      function make_diagnostic(msg, x1, y1, x2, y2, severity, source)
+      function make_diagnostic(msg, x1, y1, x2, y2, severity, source, code)
         return {
           lnum = x1,
           col = y1,
@@ -25,23 +25,24 @@ describe('vim.diagnostic', function()
           message = msg,
           severity = severity,
           source = source,
+          code = code,
         }
       end
 
-      function make_error(msg, x1, y1, x2, y2, source)
-        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.ERROR, source)
+      function make_error(msg, x1, y1, x2, y2, source, code)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.ERROR, source, code)
       end
 
-      function make_warning(msg, x1, y1, x2, y2, source)
-        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.WARN, source)
+      function make_warning(msg, x1, y1, x2, y2, source, code)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.WARN, source, code)
       end
 
-      function make_info(msg, x1, y1, x2, y2, source)
-        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.INFO, source)
+      function make_info(msg, x1, y1, x2, y2, source, code)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.INFO, source, code)
       end
 
-      function make_hint(msg, x1, y1, x2, y2, source)
-        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.HINT, source)
+      function make_hint(msg, x1, y1, x2, y2, source, code)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.HINT, source, code)
       end
 
       function count_diagnostics(bufnr, severity, namespace)
@@ -1780,6 +1781,51 @@ end)
         pcall_err(exec_lua, [[ vim.diagnostic.open_float({ prefix = 42 }) ]]))
     end)
 
+    it('can add a suffix to diagnostics', function()
+      -- Default is to render the diagnostic error code
+      eq({'1. Syntax error [code-x]', '2. Some warning [code-y]'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3, nil, "code-x"),
+          make_warning("Some warning", 1, 1, 1, 3, nil, "code-y"),
+        }
+        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({header = false, scope = "buffer"})
+        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, nil, "code-x"),
+          make_warning("Some warning", 1, 1, 1, 3, nil, "code-y"),
+        }
+        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({header = false, scope = "buffer", suffix = ""})
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      -- Suffix is rendered on the last line of a multiline diagnostic
+      eq({'1. Syntax error', '   More context [code-x]'}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error\nMore context", 0, 1, 0, 3, nil, "code-x"),
+        }
+        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({header = false, scope = "buffer"})
+        local lines = vim.api.nvim_buf_get_lines(float_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq(".../diagnostic.lua:0: suffix: expected string|table|function, got number",
+        pcall_err(exec_lua, [[ vim.diagnostic.open_float({ suffix = 42 }) ]]))
+    end)
+
     it('works with the old signature', function()
       eq({'1. Syntax error'}, exec_lua [[
         local diagnostics = {