feat: add more capabilities to `automatic_enable`

9 files changed, 225 insertions, 69 deletions

diff --git a/README.md b/README.md
index 72f01bd..a2d3e1e 100644
--- a/README.md
+++ b/README.md
@@ -83,6 +83,42 @@ runtimepath`](https://neovim.io/doc/user/options.html#'runtimepath') before sett
 
 Refer to the [Configuration](#configuration) section for information about which settings are available.
 
+# Automatically enable installed servers
+
+`mason-lspconfig.nvim` will automatically enable installed servers for you by default.
+
+To disable this feature:
+
+```lua
+require("mason-lspconfig").setup {
+    automatic_enable = false
+}
+```
+
+To exclude certain servers from being enabled:
+
+```lua
+require("mason-lspconfig").setup {
+    automatic_enable = {
+        exclude = {
+            "rust_analyzer",
+            "ts_ls"
+            }
+    }
+}
+```
+
+Alternatively, to only enable specific servers:
+
+```lua
+require("mason-lspconfig").setup {
+    automatic_enable = {
+        "lua_ls",
+        "vimls"
+    }
+}
+```
+
 # Commands
 
 > `:h mason-lspconfig-commands`
@@ -114,6 +150,22 @@ local DEFAULT_SETTINGS = {
     ensure_installed = {},
 
     -- Whether installed servers should automatically be enabled via `:h vim.lsp.enable()`.
-    automatic_enable = true
+    --
+    -- To exclude certain servers from being automatically enabled:
+    -- ```lua
+    --   automatic_enable = {
+    --     exclude = { "rust_analyzer", "ts_ls" }
+    --   }
+    -- ```
+    --
+    -- To only enable certain servers to be automatically enabled:
+    -- ```lua
+    --   automatic_enable = {
+    --     "lua_ls",
+    --     "vimls"
+    --   }
+    -- ```
+    ---@type boolean | string[] | { exclude: string[] }
+    automatic_enable = true,
 }
 ```
diff --git a/doc/mason-lspconfig.txt b/doc/mason-lspconfig.txt
index e617361..f5b309f 100644
--- a/doc/mason-lspconfig.txt
+++ b/doc/mason-lspconfig.txt
@@ -121,7 +121,23 @@ Example:
         ensure_installed = {},
 
         -- Whether installed servers should automatically be enabled via `:h vim.lsp.enable()`.
-        automatic_enable = true
+        --
+        -- To exclude certain servers from being automatically enabled:
+        -- ```lua
+        --   automatic_enable = {
+        --     exclude = { "rust_analyzer", "ts_ls" }
+        --   }
+        -- ```
+        --
+        -- To only enable certain servers to be automatically enabled:
+        -- ```lua
+        --   automatic_enable = {
+        --     "lua_ls",
+        --     "vimls"
+        --   }
+        -- ```
+        ---@type boolean | string[] | { exclude: string[] }
+        automatic_enable = true,
     }
 <
 
diff --git a/lua/mason-lspconfig/automatic_enable.lua b/lua/mason-lspconfig/automatic_enable.lua
deleted file mode 100644
index b1f0046..0000000
--- a/lua/mason-lspconfig/automatic_enable.lua
+++ /dev/null
@@ -1,32 +0,0 @@
-local _ = require "mason-core.functional"
-local mappings = require "mason-lspconfig.mappings"
-local registry = require "mason-registry"
-local settings = require "mason-lspconfig.settings"
-
----@param mason_pkg string
-local function setup_server(mason_pkg)
-    local lspconfig_name = mappings.get_mason_map().package_to_lspconfig[mason_pkg]
-    if not lspconfig_name then
-        return
-    end
-
-    -- We don't provide LSP configurations in the lsp/ directory because it risks overriding configurations in a way the
-    -- user doesn't want. Instead we only override LSP configurations for servers that are installed via Mason.
-    local ok, config = pcall(require, ("mason-lspconfig.lsp.%s"):format(lspconfig_name))
-    if ok then
-        vim.lsp.config(lspconfig_name, config)
-    end
-
-    if settings.current.automatic_enable then
-        vim.lsp.enable(lspconfig_name)
-    end
-end
-
-_.each(setup_server, registry.get_installed_package_names())
-
-registry:on(
-    "package:install:success",
-    vim.schedule_wrap(function(pkg)
-        setup_server(pkg.name)
-    end)
-)
diff --git a/lua/mason-lspconfig/features/automatic_enable.lua b/lua/mason-lspconfig/features/automatic_enable.lua
new file mode 100644
index 0000000..34ae265
--- /dev/null
+++ b/lua/mason-lspconfig/features/automatic_enable.lua
@@ -0,0 +1,54 @@
+local _ = require "mason-core.functional"
+local mappings = require "mason-lspconfig.mappings"
+local registry = require "mason-registry"
+local settings = require "mason-lspconfig.settings"
+
+---@param mason_pkg string | Package
+local function enable_server(mason_pkg)
+    if type(mason_pkg) ~= "string" then
+        mason_pkg = mason_pkg.name
+    end
+    local lspconfig_name = mappings.get_mason_map().package_to_lspconfig[mason_pkg]
+    if not lspconfig_name then
+        return
+    end
+
+    local automatic_enable = settings.current.automatic_enable
+
+    if type(automatic_enable) == "table" then
+        local exclude = automatic_enable.exclude
+        if exclude then
+            if _.any(_.equals(lspconfig_name), exclude) then
+                -- This server is explicitly excluded.
+                return
+            end
+        else
+            if not _.any(_.equals(lspconfig_name), automatic_enable) then
+                -- This server is not explicitly enabled.
+                return
+            end
+        end
+    elseif automatic_enable == false then
+        return
+    end
+
+    -- We don't provide LSP configurations in the lsp/ directory because it risks overriding configurations in a way the
+    -- user doesn't want. Instead we only override LSP configurations for servers that are installed via Mason.
+    local ok, config = pcall(require, ("mason-lspconfig.lsp.%s"):format(lspconfig_name))
+    if ok then
+        vim.lsp.config(lspconfig_name, config)
+    end
+
+    vim.lsp.enable(lspconfig_name)
+end
+
+local enable_server_scheduled = vim.schedule_wrap(enable_server)
+
+return function()
+    _.each(enable_server, registry.get_installed_package_names())
+
+    -- We deregister the event handler primarily for testing purposes where .setup() is called multiple times in the
+    -- same instance
+    registry:off("package:install:success", enable_server_scheduled)
+    registry:on("package:install:success", enable_server_scheduled)
+end
diff --git a/lua/mason-lspconfig/ensure_installed.lua b/lua/mason-lspconfig/features/ensure_installed.lua
index 583feec..583feec 100644
--- a/lua/mason-lspconfig/ensure_installed.lua
+++ b/lua/mason-lspconfig/features/ensure_installed.lua
diff --git a/lua/mason-lspconfig/init.lua b/lua/mason-lspconfig/init.lua
index ff37463..631a78f 100644
--- a/lua/mason-lspconfig/init.lua
+++ b/lua/mason-lspconfig/init.lua
@@ -27,9 +27,9 @@ function M.setup(config)
     local registry = require "mason-registry"
     registry.refresh(vim.schedule_wrap(function()
         if not platform.is_headless and #settings.current.ensure_installed > 0 then
-            require "mason-lspconfig.ensure_installed"()
+            require "lua.mason-lspconfig.features.ensure_installed"()
         end
-        require "mason-lspconfig.automatic_enable"
+        require "lua.mason-lspconfig.features.automatic_enable"()
         registry.register_package_aliases(_.map(function(server_name)
             return { server_name }
         end, require("mason-lspconfig.mappings").get_mason_map().package_to_lspconfig))
diff --git a/lua/mason-lspconfig/mappings.lua b/lua/mason-lspconfig/mappings.lua
index 98cc8c4..2523aed 100644
--- a/lua/mason-lspconfig/mappings.lua
+++ b/lua/mason-lspconfig/mappings.lua
@@ -4,13 +4,6 @@ local registry = require "mason-registry"
 local M = {}
 
 function M.get_mason_map()
-    if not registry.get_all_package_specs then
-        return {
-            lspconfig_to_package = {},
-            package_to_lspconfig = {},
-        }
-    end
-
     ---@type table<string, string>
     local package_to_lspconfig = {}
     for _, pkg_spec in ipairs(registry.get_all_package_specs()) do
diff --git a/lua/mason-lspconfig/settings.lua b/lua/mason-lspconfig/settings.lua
index cadc31e..1c6b3ec 100644
--- a/lua/mason-lspconfig/settings.lua
+++ b/lua/mason-lspconfig/settings.lua
@@ -7,6 +7,22 @@ local DEFAULT_SETTINGS = {
     ensure_installed = {},
 
     -- Whether installed servers should automatically be enabled via `:h vim.lsp.enable()`.
+    --
+    -- To exclude certain servers from being automatically enabled:
+    -- ```lua
+    --   automatic_enable = {
+    --     exclude = { "rust_analyzer", "ts_ls" }
+    --   }
+    -- ```
+    --
+    -- To only enable certain servers to be automatically enabled:
+    -- ```lua
+    --   automatic_enable = {
+    --     "lua_ls",
+    --     "vimls"
+    --   }
+    -- ```
+    ---@type boolean | string[] | { exclude: string[] }
     automatic_enable = true,
 }
 
diff --git a/tests/mason-lspconfig/setup_spec.lua b/tests/mason-lspconfig/setup_spec.lua
index 88a0a28..cb8c70d 100644
--- a/tests/mason-lspconfig/setup_spec.lua
+++ b/tests/mason-lspconfig/setup_spec.lua
@@ -10,14 +10,7 @@ local registry = require "mason-registry"
 
 describe("mason-lspconfig setup", function()
     before_each(function()
-        local settings = require "mason-lspconfig.settings"
-        settings.set(settings._DEFAULT_SETTINGS)
-
-        for _, pkg in ipairs(registry.get_all_packages()) do
-            if pkg:is_installed() then
-                pkg:uninstall()
-            end
-        end
+        a.run_blocking(a.wait, vim.schedule)
     end)
 
     it("should set up user commands", function()
@@ -42,6 +35,37 @@ describe("mason-lspconfig setup", function()
     end)
 
     it(
+        "should set up package aliases",
+        async_test(function()
+            spy.on(registry, "register_package_aliases")
+
+            mason_lspconfig.setup {}
+            a.wait(vim.schedule)
+
+            assert.spy(registry.register_package_aliases).was_called(1)
+            assert.spy(registry.register_package_aliases).was_called_with {
+                ["dummy"] = { "dummylsp" },
+                ["dummy2"] = { "dummy2lsp" },
+                ["fail_dummy"] = { "fail_dummylsp" },
+            }
+        end)
+    )
+end)
+
+describe("mason-lspconfig.setup() :: feature :: ensure_installed", function()
+    before_each(function()
+        a.run_blocking(a.wait, vim.schedule)
+        local settings = require "mason-lspconfig.settings"
+        settings.set(settings._DEFAULT_SETTINGS)
+
+        for _, pkg in ipairs(registry.get_all_packages()) do
+            if pkg:is_installed() then
+                pkg:uninstall()
+            end
+        end
+    end)
+
+    it(
         "should install servers listed in ensure_installed",
         async_test(function()
             local dummy = registry.get_package "dummy"
@@ -113,32 +137,65 @@ describe("mason-lspconfig setup", function()
             end)
         end)
     )
+end)
+
+describe("mason-lspconfig.setup() :: feature :: automatic_enable", function()
+    before_each(function()
+        a.run_blocking(a.wait, vim.schedule)
+        local settings = require "mason-lspconfig.settings"
+        settings.set(settings._DEFAULT_SETTINGS)
+
+        spy.on(vim.lsp, "enable")
+        stub(registry, "get_installed_package_names").returns {
+            "dummy",
+            "dummy2",
+        }
+    end)
 
     it(
-        "should set up package aliases",
+        "should enable all installed servers",
         async_test(function()
-            stub(registry, "register_package_aliases")
+            mason_lspconfig.setup {
+                automatic_enable = true,
+            }
 
-            local mapping_mock = stub(
-                require "mason-lspconfig.mappings",
-                "get_mason_map",
-                mockx.returns {
-                    package_to_lspconfig = {
-                        ["rust-analyzer"] = "rust_analyzer",
-                        ["typescript-language-server"] = "ts_ls",
-                    },
-                }
-            )
+            a.wait(vim.schedule)
+
+            assert.spy(vim.lsp.enable).was_called(2)
+            assert.spy(vim.lsp.enable).was_called_with "dummylsp"
+            assert.spy(vim.lsp.enable).was_called_with "dummy2lsp"
+        end)
+    )
+
+    it(
+        "should exclude servers",
+        async_test(function()
+            mason_lspconfig.setup {
+                automatic_enable = {
+                    exclude = { "dummy2lsp" },
+                },
+            }
 
-            mason_lspconfig.setup {}
             a.wait(vim.schedule)
 
-            assert.spy(registry.register_package_aliases).was_called(1)
-            assert.spy(registry.register_package_aliases).was_called_with {
-                ["rust-analyzer"] = { "rust_analyzer" },
-                ["typescript-language-server"] = { "ts_ls" },
+            assert.spy(vim.lsp.enable).was_called(1)
+            assert.spy(vim.lsp.enable).was_called_with "dummylsp"
+        end)
+    )
+
+    it(
+        "should only enable specified servers",
+        async_test(function()
+            mason_lspconfig.setup {
+                automatic_enable = {
+                    "dummy2lsp",
+                },
             }
-            mapping_mock:revert()
+
+            a.wait(vim.schedule)
+
+            assert.spy(vim.lsp.enable).was_called(1)
+            assert.spy(vim.lsp.enable).was_called_with "dummy2lsp"
         end)
     )
 end)