feat(diagnostic): allow customized diagnostic messages (#15742)

gpanders · web-flow · commit d999c96cf314 · 2021-09-22T12:20:15.000-07:00
Provide a 'format' option for virtual text and floating window previews
that allows the displayed text of a diagnostic to be customized.
diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
@@ -230,6 +230,18 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
                                    • source: (string) Include the diagnostic
                                      source in virtual text. 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. Example:>
+
+                                       function(diagnostic)
+                                         if diagnostic.severity == vim.diagnostic.severity.ERROR then
+                                           return string.format("E: %s", diagnostic.message)
+                                         end
+                                         return diagnostic.message
+                                       end
+<
 
                                  • signs: (default true) Use signs for
                                    diagnostics. Options:
@@ -519,6 +531,10 @@ show_position_diagnostics({opts}, {bufnr}, {position})
                                 • source: (string) Include the diagnostic
                                   source in 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.
                     {bufnr}     number|nil Buffer number. Defaults to the
                                 current buffer.
                     {position}  table|nil The (0,0)-indexed position. Defaults
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
@@ -76,6 +76,20 @@ local function prefix_source(source, diagnostics)
   end, diagnostics)
 end
 
+---@private
+local function reformat_diagnostics(format, diagnostics)
+  vim.validate {
+    format = {format, 'f'},
+    diagnostics = {diagnostics, 't'},
+  }
+
+  local formatted = vim.deepcopy(diagnostics)
+  for _, diagnostic in ipairs(formatted) do
+    diagnostic.message = format(diagnostic)
+  end
+  return formatted
+end
+
 ---@private
 local function resolve_optional_value(option, namespace, bufnr)
   local enabled_val = {}
@@ -376,6 +390,10 @@ local function show_diagnostics(opts, diagnostics)
     table.insert(highlights, {0, "Bold"})
   end
 
+  if opts.format then
+    diagnostics = reformat_diagnostics(opts.format, diagnostics)
+  end
+
   if opts.source then
     diagnostics = prefix_source(opts.source, diagnostics)
   end
@@ -524,6 +542,17 @@ end
 ---                       severity |diagnostic-severity|
 ---                       * source: (string) Include the diagnostic source in virtual
 ---                       text. 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. Example:
+---                       <pre>
+---                       function(diagnostic)
+---                         if diagnostic.severity == vim.diagnostic.severity.ERROR then
+---                           return string.format("E: %s", diagnostic.message)
+---                         end
+---                         return diagnostic.message
+---                       end
+---                       </pre>
 ---       - signs: (default true) Use signs for diagnostics. Options:
 ---                * severity: Only show signs for diagnostics matching the given severity
 ---                |diagnostic-severity|
@@ -865,6 +894,10 @@ function M._set_virtual_text(namespace, bufnr, diagnostics, opts)
   bufnr = get_bufnr(bufnr)
   opts = get_resolved_options({ virtual_text = opts }, namespace, bufnr).virtual_text
 
+  if opts and opts.format then
+    diagnostics = reformat_diagnostics(opts.format, diagnostics)
+  end
+
   if opts and opts.source then
     diagnostics = prefix_source(opts.source, diagnostics)
   end
@@ -1049,7 +1082,9 @@ end
 ---            - severity: See |diagnostic-severity|.
 ---            - show_header: (boolean, default true) Show "Diagnostics:" header
 ---            - source: (string) Include the diagnostic source in
----            the message. One of "always" or "if_many".
+---                      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.
 ---@param bufnr number|nil Buffer number. Defaults to the current buffer.
 ---@param position table|nil The (0,0)-indexed position. Defaults to the current cursor position.
 ---@return tuple ({popup_bufnr}, {win_id})
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
@@ -631,6 +631,61 @@ describe('vim.diagnostic', function()
       eq(' source x: Some error', result[1])
       eq(' source y: Another error', result[2])
     end)
+
+    it('supports a format function for diagnostic messages', function()
+      local result = exec_lua [[
+        vim.diagnostic.config({
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            format = function(diagnostic)
+              if diagnostic.severity == vim.diagnostic.severity.ERROR then
+                return string.format("🔥 %s", diagnostic.message)
+              end
+              return string.format("👀 %s", diagnostic.message)
+            end,
+          }
+        })
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_warning('Warning', 0, 0, 0, 0),
+          make_error('Error', 1, 0, 1, 0),
+        })
+
+        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        return {extmarks[1][4].virt_text, extmarks[2][4].virt_text}
+      ]]
+      eq(" 👀 Warning", result[1][2][1])
+      eq(" 🔥 Error", result[2][2][1])
+    end)
+
+    it('includes source for formatted diagnostics', function()
+      local result = exec_lua [[
+        vim.diagnostic.config({
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            source = 'always',
+            format = function(diagnostic)
+              if diagnostic.severity == vim.diagnostic.severity.ERROR then
+                return string.format("🔥 %s", diagnostic.message)
+              end
+              return string.format("👀 %s", diagnostic.message)
+            end,
+          }
+        })
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_warning('Warning', 0, 0, 0, 0, 'some_linter'),
+          make_error('Error', 1, 0, 1, 0, 'another_linter'),
+        })
+
+        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        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])
+    end)
   end)
 
   describe('set()', function()
