feat(diagnostic): add `suffix` option to `virt_text` config (#21140)

This introduces a `suffix` option to the `virt_text` config in
`vim.diagnostic.config()`. The suffix can either be a string which is appended
to the diagnostic message or a function returning such. The function receives a
`diagnostic` argument, which is the diagnostic table of the last diagnostic (the
one whose message is rendered as virt text).

3 files changed, 52 insertions, 1 deletions

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index 53046b399f..a2bacebc72 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -365,6 +365,11 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
                          the beginning of the virtual text.
                        • prefix: (string) Prepend diagnostic message with
                          prefix.
+                       • suffix: (string or function) Append diagnostic
+                         message with suffix. If a function, it must have the
+                         signature (diagnostic) -> string, where {diagnostic}
+                         is of type |diagnostic-structure|. This can be used
+                         to render an LSP diagnostic error code.
                        • 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. Example: >
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 008b00a9a0..a7fc47a5a8 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -613,6 +613,10 @@ end
 ---                       * spacing: (number) Amount of empty spaces inserted at the beginning
 ---                                  of the virtual text.
 ---                       * prefix: (string) Prepend diagnostic message with prefix.
+---                       * suffix: (string or function) Append diagnostic message with suffix.
+---                                 If a function, it must have the signature (diagnostic) ->
+---                                 string, where {diagnostic} is of type |diagnostic-structure|.
+---                                 This can be used to render an LSP diagnostic error code.
 ---                       * 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. Example:
@@ -1039,6 +1043,7 @@ function M._get_virt_text_chunks(line_diags, opts)
 
   opts = opts or {}
   local prefix = opts.prefix or '■'
+  local suffix = opts.suffix or ''
   local spacing = opts.spacing or 4
 
   -- Create a little more space between virtual text and contents
@@ -1052,8 +1057,11 @@ function M._get_virt_text_chunks(line_diags, opts)
   -- TODO(tjdevries): Allow different servers to be shown first somehow?
   -- TODO(tjdevries): Display server name associated with these?
   if last.message then
+    if type(suffix) == 'function' then
+      suffix = suffix(last) or ''
+    end
     table.insert(virt_texts, {
-      string.format('%s %s', prefix, last.message:gsub('\r', ''):gsub('\n', '  ')),
+      string.format('%s %s%s', prefix, last.message:gsub('\r', ''):gsub('\n', '  '), suffix),
       virtual_text_highlight_map[last.severity],
     })
 
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 816b359725..b7fe39cc91 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -1161,6 +1161,44 @@ end)
       eq(" some_linter: 👀 Warning", result[1][2][1])
       eq(" another_linter: 🔥 Error", result[2][2][1])
     end)
+
+    it('can add a suffix to virtual text', function()
+      eq(' Some error ✘',  exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            suffix = ' ✘',
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local virt_text = extmarks[1][4].virt_text[2][1]
+        return virt_text
+      ]])
+
+      eq(' Some error [err-code]',  exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0, nil, 'err-code'),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            suffix = function(diag) return string.format(' [%s]', diag.code) end,
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local virt_text = extmarks[1][4].virt_text[2][1]
+        return virt_text
+      ]])
+    end)
   end)
 
   describe('set()', function()