From 657df23a2c5811b593b00202fec4f77014b058b6 Mon Sep 17 00:00:00 2001
From: Scott Penrose <penrose@gmail.com>
Date: Sat, 29 Jun 2024 14:35:52 -0400
Subject: [PATCH] feat: add trouble diagnostic counts to statusline

Allows you to use trouble diagnostic counts on your status line. This is
best used when you are running a filter on trouble diagnostics. For
example, I only want to see workspace diagnostics and counts.
---
 README.md                 | 37 ++++++++++++++++++++++++++++-
 lua/trouble/api.lua       | 50 +++++++++++++++++++++++++++++++++++++++
 lua/trouble/view/text.lua | 20 ++++++++++++++++
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 5d205a42..f733d497 100644
--- a/README.md
+++ b/README.md
@@ -637,7 +637,7 @@ config.defaults.actions.files["ctrl-t"] = actions.open
 
 When you open fzf-lua, you can now hit `<c-t>` to open the results in **Trouble**
 
-### Statusline Component
+### Statusline LSP Document Symbols Component
 
 Example for [lualine.nvim](https://github.com/nvim-lualine/lualine.nvim):
 
@@ -664,6 +664,41 @@ Example for [lualine.nvim](https://github.com/nvim-lualine/lualine.nvim):
 }
 ```
 
+### Statusline Diagnostic Count Component
+
+Have the diagnostics counts in statusline widget match the trouble diagnostics filter you use.
+
+Example for [lualine.nvim](https://github.com/nvim-lualine/lualine.nvim):
+
+```lua
+{
+  "nvim-lualine/lualine.nvim",
+  opts = function(_, opts)
+    local trouble = require 'trouble'
+    local troubleDignosticsCount = trouble.diagnosticsCount ({
+      -- use diagnostics mode and specify filter
+      mode = 'diagnostics',
+      filter = {
+        -- limit to files in the current project
+        function(item)
+          return item.filename:find((vim.loop or vim.uv).cwd(), 1, true)
+        end,
+      },
+      -- Or use a custom mode you created that already contains the filter
+      -- mode = 'onlyworkspace',
+    })
+    table.insert(opts.sections.lualine_c, {
+      'diagnostics',
+      sources = {
+        function()
+          return troubleDignosticsCount.get()
+        end,
+      },
+    })
+  end,
+}
+```
+
 ## 🎨 Colors
 
 The table below shows all the highlight groups defined for Trouble.
diff --git a/lua/trouble/api.lua b/lua/trouble/api.lua
index 8bb5bbd3..9df17041 100644
--- a/lua/trouble/api.lua
+++ b/lua/trouble/api.lua
@@ -193,6 +193,56 @@ function M.statusline(opts)
   }
 end
 
+-- Renders a trouble trable of diagnostic counts
+--   { error=error_cnt, warn=warn_cnt, info=info_cnt, hint=hint_cnt }
+-- Check the docs for examples.
+---@param opts? trouble.Mode|string|{hl_group?:string}
+---@return {get: (fun():table)}
+function M.diagnosticsCount(opts)
+  local Spec = require("trouble.spec")
+  local Section = require("trouble.view.section")
+  local Render = require("trouble.view.render")
+  opts.groups = {
+    { "severity", format = "{severity}{count}" },
+  }
+  opts.title = false
+  opts.format = ""
+  opts = Config.get(opts)
+
+  local renderer = Render.new(opts, {
+    multiline = false,
+    indent = false,
+  })
+  local status = nil ---@type table?
+  ---@cast opts trouble.Mode
+
+  local s = Spec.section(opts)
+  s.max_items = s.max_items or opts.max_items
+  local section = Section.new(s, opts)
+  section.on_update = function()
+    status = nil
+    if package.loaded["lualine"] then
+      vim.schedule(function()
+        require("lualine").refresh()
+      end)
+    else
+      vim.cmd.redrawstatus()
+    end
+  end
+  section:listen()
+  section:refresh()
+  return {
+    get = function()
+      if status then
+        return status
+      end
+      renderer:clear()
+      renderer:sections({ section })
+      return renderer:diagnosticCount()
+    end,
+  }
+end
+
 return setmetatable(M, {
   __index = function(_, k)
     if k == "last_mode" then
diff --git a/lua/trouble/view/text.lua b/lua/trouble/view/text.lua
index a6632c2b..d04a3b8c 100644
--- a/lua/trouble/view/text.lua
+++ b/lua/trouble/view/text.lua
@@ -122,6 +122,26 @@ function M:statusline(opts)
   return table.concat(lines, sep)
 end
 
+---@return table -- { error = count, warn = count, info = count, hint = count }
+function M:diagnosticCount()
+  local list = {}
+  for _, line in ipairs(self._lines) do
+    for _, segment in ipairs(line) do
+      local str = segment.str:gsub("%s+", "")
+      table.insert(list, str)
+    end
+  end
+
+  local lookupTable = {}
+  for i = 1, #list, 2 do
+    local key = list[i]:lower()
+    local value = tonumber(list[i + 1])
+    lookupTable[key] = value
+  end
+
+  return lookupTable
+end
+
 function M:render(buf)
   local lines = {}