feat(diagnostic): specify diagnostic virtual text prefix as a function

- vim.diagnostic.config() now accepts a function for the virtual_text.prefix
  option, which allows for rendering e.g., diagnostic severities differently.

4 files changed, 72 insertions, 16 deletions

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index 025a7dccfc..bfcc1c9092 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -392,8 +392,11 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
                          always show the diagnostic source.
                        • spacing: (number) Amount of empty spaces inserted at
                          the beginning of the virtual text.
-                       • prefix: (string) Prepend diagnostic message with
-                         prefix.
+                       • prefix: (string or function) prepend diagnostic
+                         message with prefix. If a function, it must have the
+                         signature (diagnostic) -> string, where {diagnostic}
+                         is of type |diagnostic-structure|. This can be used
+                         to render diagnostic symbols or error codes.
                        • suffix: (string or function) Append diagnostic
                          message with suffix. If a function, it must have the
                          signature (diagnostic) -> string, where {diagnostic}
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index f196c3d445..dc541fcf37 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -48,6 +48,9 @@ The following changes to existing APIs or features add new behavior.
 
 • |vim.region()| can use a string accepted by |getpos()| as position.
 
+• vim.diagnostic.config() now accepts a function for the virtual_text.prefix
+  option, which allows for rendering e.g., diagnostic severities differently.
+
 ==============================================================================
 REMOVED FEATURES                                                 *news-removed*
 
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index d1b50304c7..0d1d01b391 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -600,7 +600,10 @@ end
 ---                                 means to always show the diagnostic source.
 ---                       * spacing: (number) Amount of empty spaces inserted at the beginning
 ---                                  of the virtual text.
----                       * prefix: (string) Prepend diagnostic message with prefix.
+---                       * prefix: (string or function) prepend diagnostic message with prefix.
+---                                 If a function, it must have the signature (diagnostic) -> string,
+---                                 where {diagnostic} is of type |diagnostic-structure|. This can be
+---                                 used to render diagnostic symbols or error codes.
 ---                       * 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|.
@@ -1066,8 +1069,15 @@ function M._get_virt_text_chunks(line_diags, opts)
   -- Create a little more space between virtual text and contents
   local virt_texts = { { string.rep(' ', spacing) } }
 
-  for i = 1, #line_diags - 1 do
-    table.insert(virt_texts, { prefix, virtual_text_highlight_map[line_diags[i].severity] })
+  for i = 1, #line_diags do
+    local resolved_prefix = prefix
+    if type(prefix) == 'function' then
+      resolved_prefix = prefix(line_diags[i]) or ''
+    end
+    table.insert(
+      virt_texts,
+      { resolved_prefix, virtual_text_highlight_map[line_diags[i].severity] }
+    )
   end
   local last = line_diags[#line_diags]
 
@@ -1078,7 +1088,7 @@ function M._get_virt_text_chunks(line_diags, opts)
       suffix = suffix(last) or ''
     end
     table.insert(virt_texts, {
-      string.format('%s %s%s', prefix, last.message:gsub('\r', ''):gsub('\n', '  '), suffix),
+      string.format(' %s%s', 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 6b4ca5ac73..17b2d7da4f 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -1042,7 +1042,7 @@ end)
           local virt_text = get_virt_text_extmarks(diagnostic_ns)[1][4].virt_text
 
           local virt_texts = {}
-          for i = 2, #virt_text do
+          for i = 2, #virt_text - 1 do
             table.insert(virt_texts, (string.gsub(virt_text[i][2], "DiagnosticVirtualText", "")))
           end
 
@@ -1086,7 +1086,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]]
       eq(' source x: Some error', result)
@@ -1101,7 +1101,7 @@ end)
         }, diagnostic_ns)
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]]
       eq(' Some error', result)
@@ -1121,7 +1121,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = {extmarks[1][4].virt_text[2][1], extmarks[2][4].virt_text[2][1]}
+        local virt_text = {extmarks[1][4].virt_text[3][1], extmarks[2][4].virt_text[3][1]}
         return virt_text
       ]]
       eq(' source x: Some error', result[1])
@@ -1151,8 +1151,8 @@ end)
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
         return {extmarks[1][4].virt_text, extmarks[2][4].virt_text}
       ]]
-      eq(" 👀 Warning", result[1][2][1])
-      eq(" 🔥 Error", result[2][2][1])
+      eq(" 👀 Warning", result[1][3][1])
+      eq(" 🔥 Error", result[2][3][1])
     end)
 
     it('includes source for formatted diagnostics', function()
@@ -1179,8 +1179,48 @@ end)
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
         return {extmarks[1][4].virt_text, extmarks[2][4].virt_text}
       ]]
-      eq(" some_linter: 👀 Warning", result[1][2][1])
-      eq(" another_linter: 🔥 Error", result[2][2][1])
+      eq(" some_linter: 👀 Warning", result[1][3][1])
+      eq(" another_linter: 🔥 Error", result[2][3][1])
+    end)
+
+    it('can add a prefix to virtual text', function()
+      eq('E 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 = 'E',
+            suffix = '',
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local prefix = extmarks[1][4].virt_text[2][1]
+        local message = extmarks[1][4].virt_text[3][1]
+        return prefix .. message
+      ]])
+
+      eq('[err-code] Some error',  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 = function(diag) return string.format('[%s]', diag.code) end,
+            suffix = '',
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local prefix = extmarks[1][4].virt_text[2][1]
+        local message = extmarks[1][4].virt_text[3][1]
+        return prefix .. message
+      ]])
     end)
 
     it('can add a suffix to virtual text', function()
@@ -1198,7 +1238,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]])
 
@@ -1216,7 +1256,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]])
     end)