Merge pull request #15717 from gpanders/diagnostic-include-source

feat(diagnostic): add option to include diagnostic source

5 files changed, 178 insertions, 60 deletions

diff --git a/runtime/doc/deprecated.txt b/runtime/doc/deprecated.txt
index a0c291964e..a7ce4135af 100644
--- a/runtime/doc/deprecated.txt
+++ b/runtime/doc/deprecated.txt
@@ -75,7 +75,9 @@ For each of the functions below, use the corresponding function in
 *vim.lsp.diagnostic.get_prev()*
 *vim.lsp.diagnostic.get_prev_pos()*
 *vim.lsp.diagnostic.get_virtual_text_chunks_for_line()*
-			Use |vim.diagnostic.get_virt_text_chunks()| instead.
+				No replacement. Use options provided by
+				|vim.diagnostic.config()| to customize
+				virtual text.
 *vim.lsp.diagnostic.goto_next()*
 *vim.lsp.diagnostic.goto_prev()*
 *vim.lsp.diagnostic.redraw()*		Use |vim.diagnostic.show()| instead.
diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index 59b73771a6..c0ae6a5a1a 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -47,6 +47,7 @@ A diagnostic is a Lua table with the following keys:
 	end_col: The final column of the diagnostic
 	severity: The severity of the diagnostic |vim.diagnostic.severity|
 	message: The diagnostic text
+	source: The source of the diagnostic
 
 Diagnostics use the same indexing as the rest of the Nvim API (i.e. 0-based
 rows and columns). |api-indexing|
@@ -226,6 +227,9 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
                                    • severity: Only show virtual text for
                                      diagnostics matching the given severity
                                      |diagnostic-severity|
+                                   • source: (string) Include the diagnostic
+                                     source in virtual text. One of "always"
+                                     or "if_many".
 
                                  • signs: (default true) Use signs for
                                    diagnostics. Options:
@@ -329,26 +333,6 @@ get_prev_pos({opts})                           *vim.diagnostic.get_prev_pos()*
                 Return: ~
                     table Previous diagnostic position as a (row, col) tuple.
 
-                                       *vim.diagnostic.get_virt_text_chunks()*
-get_virt_text_chunks({line_diags}, {opts})
-                Get virtual text chunks to display using
-                |nvim_buf_set_extmark()|.
-
-                Parameters: ~
-                    {line_diags}  table The diagnostics associated with the
-                                  line.
-                    {opts}        table|nil Configuration table with the
-                                  following keys:
-                                  • prefix: (string) Prefix to display before
-                                    virtual text on line.
-                                  • spacing: (number) Number of spaces to
-                                    insert before virtual text.
-
-                Return: ~
-                    array of ({text}, {hl_group}) tuples. This can be passed
-                    directly to the {virt_text} option of
-                    |nvim_buf_set_extmark()|.
-
 goto_next({opts})                                 *vim.diagnostic.goto_next()*
                 Move to the next diagnostic.
 
@@ -532,6 +516,9 @@ show_position_diagnostics({opts}, {bufnr}, {position})
                                 • 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".
                     {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
index 0261475f72..55bf212389 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -48,6 +48,35 @@ local function filter_by_severity(severity, diagnostics)
 end
 
 ---@private
+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'" } }
+
+  if source == "if_many" then
+    local sources = {}
+    for _, d in pairs(diagnostics) do
+      if d.source then
+        sources[d.source] = true
+      end
+    end
+    if #vim.tbl_keys(sources) <= 1 then
+      return diagnostics
+    end
+  end
+
+  return vim.tbl_map(function(d)
+    if not d.source then
+      return d
+    end
+
+    local t = vim.deepcopy(d)
+    t.message = string.format("%s: %s", d.source, d.message)
+    return t
+  end, diagnostics)
+end
+
+---@private
 local function resolve_optional_value(option, namespace, bufnr)
   local enabled_val = {}
 
@@ -336,7 +365,9 @@ end
 ---@param diagnostics table: The diagnostics to display
 ---@return table {popup_bufnr, win_id}
 local function show_diagnostics(opts, diagnostics)
-  if vim.tbl_isempty(diagnostics) then return end
+  if vim.tbl_isempty(diagnostics) then
+    return
+  end
   local lines = {}
   local highlights = {}
   local show_header = vim.F.if_nil(opts.show_header, true)
@@ -345,6 +376,10 @@ local function show_diagnostics(opts, diagnostics)
     table.insert(highlights, {0, "Bold"})
   end
 
+  if opts.source then
+    diagnostics = prefix_source(opts.source, diagnostics)
+  end
+
   for i, diagnostic in ipairs(diagnostics) do
     local prefix = string.format("%d. ", i)
     local hiname = floating_highlight_map[diagnostic.severity]
@@ -487,6 +522,8 @@ end
 ---       - virtual_text: (default true) Use virtual text for diagnostics. Options:
 ---                       * severity: Only show virtual text for diagnostics matching the given
 ---                       severity |diagnostic-severity|
+---                       * source: (string) Include the diagnostic source in virtual
+---                       text. One of "always" or "if_many".
 ---       - signs: (default true) Use signs for diagnostics. Options:
 ---                * severity: Only show signs for diagnostics matching the given severity
 ---                |diagnostic-severity|
@@ -814,6 +851,8 @@ end
 ---@param opts table|nil Configuration table with the following keys:
 ---            - prefix: (string) Prefix to display before virtual text on line.
 ---            - spacing: (number) Number of spaces to insert before virtual text.
+---            - source: (string) Include the diagnostic source in virtual text. One of "always" or
+---                      "if_many".
 ---@private
 function M._set_virtual_text(namespace, bufnr, diagnostics, opts)
   vim.validate {
@@ -826,12 +865,16 @@ 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.source then
+    diagnostics = prefix_source(opts.source, diagnostics)
+  end
+
   local buffer_line_diagnostics = diagnostic_lines(diagnostics)
   for line, line_diagnostics in pairs(buffer_line_diagnostics) do
     if opts and opts.severity then
       line_diagnostics = filter_by_severity(opts.severity, line_diagnostics)
     end
-    local virt_texts = M.get_virt_text_chunks(line_diagnostics, opts)
+    local virt_texts = M._get_virt_text_chunks(line_diagnostics, opts)
 
     if virt_texts then
       vim.api.nvim_buf_set_extmark(bufnr, namespace, line, 0, {
@@ -844,13 +887,11 @@ end
 
 --- Get virtual text chunks to display using |nvim_buf_set_extmark()|.
 ---
----@param line_diags table The diagnostics associated with the line.
----@param opts table|nil Configuration table with the following keys:
----            - prefix: (string) Prefix to display before virtual text on line.
----            - spacing: (number) Number of spaces to insert before virtual text.
----@return array of ({text}, {hl_group}) tuples. This can be passed directly to
----        the {virt_text} option of |nvim_buf_set_extmark()|.
-function M.get_virt_text_chunks(line_diags, opts)
+--- Exported for backward compatibility with
+--- vim.lsp.diagnostic.get_virtual_text_chunks_for_line(). When that function is eventually removed,
+--- this can be made local.
+---@private
+function M._get_virt_text_chunks(line_diags, opts)
   if #line_diags == 0 then
     return nil
   end
@@ -1007,6 +1048,8 @@ end
 ---            - namespace: (number) Limit diagnostics to the given namespace
 ---            - 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".
 ---@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/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index a9d93ae65f..c6c08a15d3 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -103,6 +103,7 @@ local function diagnostic_lsp_to_vim(diagnostics, bufnr, client_id)
       end_col = line_byte_from_position(buf_lines, _end.line, _end.character, offset_encoding),
       severity = severity_lsp_to_vim(diagnostic.severity),
       message = diagnostic.message,
+      source = diagnostic.source,
       user_data = {
         lsp = {
           code = diagnostic.code,
@@ -132,6 +133,7 @@ local function diagnostic_vim_to_lsp(diagnostics)
       },
       severity = severity_vim_to_lsp(diagnostic.severity),
       message = diagnostic.message,
+      source = diagnostic.source,
     }, diagnostic.user_data and (diagnostic.user_data.lsp or {}) or {})
   end, diagnostics)
 end
@@ -527,7 +529,7 @@ end
 ---@return an array of [text, hl_group] arrays. This can be passed directly to
 ---        the {virt_text} option of |nvim_buf_set_extmark()|.
 function M.get_virtual_text_chunks_for_line(bufnr, _, line_diags, opts)
-  return vim.diagnostic.get_virt_text_chunks(diagnostic_lsp_to_vim(line_diags, bufnr), opts)
+  return vim.diagnostic._get_virt_text_chunks(diagnostic_lsp_to_vim(line_diags, bufnr), opts)
 end
 
 --- Open a floating window with the diagnostics from {position}
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index e542b02abe..7c5a52827c 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -14,48 +14,32 @@ describe('vim.diagnostic', function()
     exec_lua [[
       require('vim.diagnostic')
 
-      function make_error(msg, x1, y1, x2, y2)
+      function make_diagnostic(msg, x1, y1, x2, y2, severity, source)
         return {
           lnum = x1,
           col = y1,
           end_lnum = x2,
           end_col = y2,
           message = msg,
-          severity = vim.diagnostic.severity.ERROR,
+          severity = severity,
+          source = source,
         }
       end
 
-      function make_warning(msg, x1, y1, x2, y2)
-        return {
-          lnum = x1,
-          col = y1,
-          end_lnum = x2,
-          end_col = y2,
-          message = msg,
-          severity = vim.diagnostic.severity.WARN,
-        }
+      function make_error(msg, x1, y1, x2, y2, source)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.ERROR, source)
       end
 
-      function make_info(msg, x1, y1, x2, y2)
-        return {
-          lnum = x1,
-          col = y1,
-          end_lnum = x2,
-          end_col = y2,
-          message = msg,
-          severity = vim.diagnostic.severity.INFO,
-        }
+      function make_warning(msg, x1, y1, x2, y2, source)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.WARN, source)
       end
 
-      function make_hint(msg, x1, y1, x2, y2)
-        return {
-          lnum = x1,
-          col = y1,
-          end_lnum = x2,
-          end_col = y2,
-          message = msg,
-          severity = vim.diagnostic.severity.HINT,
-        }
+      function make_info(msg, x1, y1, x2, y2, source)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.INFO, source)
+      end
+
+      function make_hint(msg, x1, y1, x2, y2, source)
+        return make_diagnostic(msg, x1, y1, x2, y2, vim.diagnostic.severity.HINT, source)
       end
 
       function count_diagnostics(bufnr, severity, namespace)
@@ -590,6 +574,63 @@ describe('vim.diagnostic', function()
       eq({'Error', 'Warn', 'Info'}, result[1])
       eq({'Info', 'Warn', 'Error'}, result[2])
     end)
+
+    it('can show diagnostic sources in virtual text', function()
+      local result = exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0, 'source x'),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            source = 'always',
+          }
+        })
+
+        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local virt_text = extmarks[1][4].virt_text[2][1]
+        return virt_text
+      ]]
+      eq(' source x: Some error', result)
+
+      result = exec_lua [[
+        vim.diagnostic.config({
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            source = 'if_many',
+          }
+        }, diagnostic_ns)
+
+        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local virt_text = extmarks[1][4].virt_text[2][1]
+        return virt_text
+      ]]
+      eq(' Some error', result)
+
+      result = exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0, 'source x'),
+          make_error('Another error', 1, 1, 1, 1, 'source y'),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = '',
+            source = 'if_many',
+          }
+        })
+
+        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local virt_text = {extmarks[1][4].virt_text[2][1], extmarks[2][4].virt_text[2][1]}
+        return virt_text
+      ]]
+      eq(' source x: Some error', result[1])
+      eq(' source y: Another error', result[2])
+    end)
   end)
 
   describe('set()', function()
@@ -854,6 +895,49 @@ describe('vim.diagnostic', function()
         return #vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
       ]])
     end)
+
+    it('can show diagnostic source', function()
+      exec_lua [[vim.api.nvim_win_set_buf(0, diagnostic_bufnr)]]
+
+      eq({"1. Syntax error"}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3, "source x"),
+        }
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local popup_bufnr, winnr = vim.diagnostic.show_line_diagnostics {
+          show_header = false,
+          source = "if_many",
+        }
+        local lines = vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq({"1. source x: Syntax error"}, exec_lua [[
+        local popup_bufnr, winnr = vim.diagnostic.show_line_diagnostics {
+          show_header = false,
+          source = "always",
+        }
+        local lines = vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+
+      eq({"1. source x: Syntax error", "2. source y: Another error"}, exec_lua [[
+        local diagnostics = {
+          make_error("Syntax error", 0, 1, 0, 3, "source x"),
+          make_error("Another error", 0, 1, 0, 3, "source y"),
+        }
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics)
+        local popup_bufnr, winnr = vim.diagnostic.show_line_diagnostics {
+          show_header = false,
+          source = "if_many",
+        }
+        local lines = vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
+        vim.api.nvim_win_close(winnr, true)
+        return lines
+      ]])
+    end)
   end)
 
   describe('set_signs()', function()