From e7801775060e2d8f9f20572fac687f438e81caa0 Mon Sep 17 00:00:00 2001
From: Michael Strobel <71396679+Kibadda@users.noreply.github.com>
Date: Wed, 16 Aug 2023 15:49:14 +0200
Subject: [PATCH] feat(diagnostic): filter diagnostics by specific severities
 (#24736)

Allow users to filter diagnostics by specifying severities
---
 runtime/doc/diagnostic.txt              | 13 +++++++++++--
 runtime/doc/news.txt                    |  3 +++
 runtime/lua/vim/diagnostic.lua          | 17 ++++++++++++++---
 test/functional/lua/diagnostic_spec.lua | 24 +++++++++++++++++++++++-
 4 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index edbba6ac41..d0aa538011 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -66,7 +66,7 @@ The "severity" key in a diagnostic is one of the values defined in
     vim.diagnostic.severity.HINT
 
 Functions that take a severity as an optional parameter (e.g.
-|vim.diagnostic.get()|) accept one of two forms:
+|vim.diagnostic.get()|) accept one of three forms:
 
 1. A single |vim.diagnostic.severity| value: >lua
 
@@ -75,8 +75,17 @@ Functions that take a severity as an optional parameter (e.g.
 2. A table with a "min" or "max" key (or both): >lua
 
     vim.diagnostic.get(0, { severity = { min = vim.diagnostic.severity.WARN } })
+<
+    This form allows users to specify a range of severities.
 
-The latter form allows users to specify a range of severities.
+3. A list-like table: >lua
+
+    vim.diagnostic.get(0, { severity = {
+        vim.diagnostic.severity.WARN,
+        vim.diagnostic.severity.INFO,
+    } })
+<
+    This form allows users to filter for specific severities
 
 ==============================================================================
 HANDLERS                                                *diagnostic-handlers*
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 576c56f408..483e829770 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -145,6 +145,9 @@ The following new APIs and features were added.
 
 • Improved messages for type errors in `vim.api.*` calls (including `opts` params)
 
+• Functions that take a severity as an optional parameter (e.g.
+  |vim.diagnostic.get()|) now also accept a list of severities |vim.diagnostic.severity|
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 1391dafd75..96d1cb7629 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -94,11 +94,22 @@ local function filter_by_severity(severity, diagnostics)
     end, diagnostics)
   end
 
-  local min_severity = to_severity(severity.min) or M.severity.HINT
-  local max_severity = to_severity(severity.max) or M.severity.ERROR
+  if severity.min or severity.max then
+    local min_severity = to_severity(severity.min) or M.severity.HINT
+    local max_severity = to_severity(severity.max) or M.severity.ERROR
+
+    return vim.tbl_filter(function(t)
+      return t.severity <= min_severity and t.severity >= max_severity
+    end, diagnostics)
+  end
+
+  local severities = {}
+  for _, s in ipairs(severity) do
+    severities[to_severity(s)] = true
+  end
 
   return vim.tbl_filter(function(t)
-    return t.severity <= min_severity and t.severity >= max_severity
+    return severities[t.severity]
   end, diagnostics)
 end
 
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 17b2d7da4f..27ba70f057 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -860,7 +860,7 @@ end)
       ]])
     end)
 
-    it('returns only requested diagnostics when severity is supplied', function()
+    it('returns only requested diagnostics when severity range is supplied', function()
       eq({2, 3, 2}, exec_lua [[
         vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
           make_error("Error 1", 1, 1, 1, 5),
@@ -882,6 +882,28 @@ end)
       ]])
     end)
 
+    it('returns only requested diagnostics when severities are supplied', function()
+      eq({1, 1, 2}, exec_lua [[
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_error("Error 1", 1, 1, 1, 5),
+          make_warning("Warning on Server 1", 1, 1, 2, 3),
+          make_info("Ignored information", 1, 1, 2, 3),
+          make_hint("Here's a hint", 1, 1, 2, 3),
+        })
+
+        return {
+          #vim.diagnostic.get(diagnostic_bufnr, { severity = {vim.diagnostic.severity.WARN} }),
+          #vim.diagnostic.get(diagnostic_bufnr, { severity = {vim.diagnostic.severity.ERROR} }),
+          #vim.diagnostic.get(diagnostic_bufnr, {
+            severity = {
+              vim.diagnostic.severity.INFO,
+              vim.diagnostic.severity.WARN,
+            }
+          }),
+        }
+      ]])
+    end)
+
     it('allows filtering by line', function()
       eq(1, exec_lua [[
         vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {