Add API based spellcheck methods

README.md

@@ -28,12 +28,13 @@ Pass any of the following options to `require("spellwarn").setup()`:
 {
     event = { -- event(s) to refresh diagnostics on
         "CursorHold",
+        "InsertLeave",
         "TextChanged",
         "TextChangedI",
         "TextChangedP",
         "TextChangedT",
     },
-    ft_config = { -- filetypes to override ft_default for
+    ft_config = { -- spellcheck method: "cursor", "iter", "treesitter", or boolean
         alpha   = false,
         help    = false,
         lazy    = false,
@@ -42,7 +43,7 @@ Pass any of the following options to `require("spellwarn").setup()`:
     },
     ft_default = true, -- whether to enable or disable for all filetypes by default
     max_file_size = nil, -- maximum file size to check in lines (nil for no limit)
-    severity = { -- severity for each spelling error type (false to disable)
+    severity = { -- severity for each spelling error type (false to disable diagnostics for that type)
         spellbad   = "WARN",
         spellcap   = "HINT",
         spelllocal = "HINT",
@@ -51,7 +52,7 @@ Pass any of the following options to `require("spellwarn").setup()`:
     prefix = "possible misspelling(s): ", -- prefix for each diagnostic message
 }
 ```
-Note that most options are overwritten (e.g. passing `ft_config = { python = false }` will mean that `alpha`, `mason`, etc. are set to true) but that `severity` is merged, so that passing `spellbad = "HINT"` won't cause `spellcap` to be nil.
+Note that most options are overwritten (e.g. passing `ft_config = { python = false }` will mean that `alpha`, `mason`, etc. are set to true) but that `severity` is merged, so that passing `spellbad = "HINT"` won't cause `spellcap` to be nil. You can pass any of `cursor`, `iter`, `treesitter`, `false`, or `true` as options to `ft_config`. The default method is `cursor`, which iterates through the buffer with `]s`. There is also `iter`, which uses the Lua API, and `treesitter`, which uses the Lua API and Treesitter (and falls back on `iter` if Treesitter is unavailable). Finally, `false` disables Spellwarn for that filetype and `true` uses the default (`cursor`).
 
 ## Usage
 The plugin should be good to go after installation with the provided snippet. It has sensible defaults. Run `:Spellwarn enable` or `:Spellwarn disable` to enable/disable during runtime (though this will *not* override `max_file_size`, `ft_config`, or `ft_default`). To disable diagnostics on a specific line, add `spellwarn:disable-next-line` to the line immediately above or `spellwarn:disable-line` to a comment at the end of the line. To disable diagnostics in a file, add a comment with `spellwarn:disable` to the *first* line of the file.

lua/spellwarn/diagnostics.lua

@@ -20,9 +20,8 @@ function M.update_diagnostics(opts, bufnr)
         return
     end
 
-    local errors = require("spellwarn.spelling").get_spelling_errors(bufnr)
     local diags = {}
-    for _, error in pairs(errors) do
+    for _, error in pairs(require("spellwarn.spelling").get_spelling_errors_main(opts, bufnr) or {}) do
         if error.word ~= "" and error.word ~= "spellwarn" then
             if opts.severity[error.type] then
                 diags[#diags + 1] = {

lua/spellwarn/init.lua

@@ -3,12 +3,13 @@ local defaults = {
     -- FIX: Trouble.nvim jump to diagnostic is slightly buggy with `TextChanged` event; no good workaround though AFAICT
     event = { -- event(s) to refresh diagnostics on
         "CursorHold",
+        "InsertLeave",
         "TextChanged",
         "TextChangedI",
         "TextChangedP",
         "TextChangedT",
     },
-    ft_config = { -- filetypes to override ft_default for
+    ft_config = { -- spellcheck method: "cursor", "iter", "treesitter", or boolean
         alpha   = false,
         help    = false,
         lazy    = false,

lua/spellwarn/spelling.lua

@@ -8,13 +8,37 @@ function M.get_error_type(word, bufnr)
     end)
 end
 
-function M.get_spelling_errors(bufnr)
+function M.check_spellwarn_comment(bufnr, linenr) -- Check for spellwarn:disable* comments
+    local above = (linenr > 1 and vim.api.nvim_buf_get_lines(bufnr, linenr - 2, linenr - 1, false)[1]) or ""
+    local above_val = string.find(above, "spellwarn:disable-next-line", 1, true) ~= nil
+    local cur = vim.api.nvim_buf_get_lines(bufnr, linenr - 1, linenr, false)[1]
+    local cur_val = string.find(cur, "spellwarn:disable-line", 1, true) ~= nil
+    return above_val or cur_val
+end
+
+function M.get_spelling_errors_main(opts, bufnr)
+    local bufopts = opts.ft_config[vim.o.filetype] or opts.ft_default
+    local disable_comment =  string.find(vim.fn.getline(1), "spellwarn:disable", 1, true) ~= nil
+
+    if vim.api.nvim_get_mode().mode == "i" or disable_comment or not bufopts then
+        return {}
+    elseif bufopts == true or bufopts == "cursor" then
+        return M.get_spelling_errors_cursor(bufnr)
+    elseif bufopts == "iter" then
+        return M.get_spelling_errors_iter(bufnr)
+    elseif bufopts == "treesitter" then
+        return M.get_spelling_errors_ts(bufnr)
+    else
+        error("Invalid value for ft_config: " .. bufopts)
+    end
+end
+
+function M.get_spelling_errors_cursor(bufnr)
     -- Save current window view and create table to store errors
     local window = vim.fn.winsaveview()
     local foldstatus = vim.o.foldenable
     local concealstatus = vim.o.conceallevel
     local errors = {}
-    if not vim.o.spell or string.find(vim.fn.getline(1), "spellwarn:disable", 1, true) ~= nil then return errors end
 
     -- Get location of first spelling error to start while loop
     vim.o.foldenable = false
@@ -24,16 +48,8 @@ function M.get_spelling_errors(bufnr)
     vim.cmd("silent normal! ]s")
     local location = vim.fn.getpos(".")
 
-    local function check_spellwarn_comment() -- Check for spellwarn:disable* comments
-        local current_line_number = vim.fn.line(".")
-        local above = (current_line_number > 1 and vim.fn.getline(current_line_number - 1)) or ""
-        local above_val = string.find(above, "spellwarn:disable-next-line", 1, true) ~= nil
-        local cur = vim.fn.getline(current_line_number)
-        local cur_val = string.find(cur, "spellwarn:disable-line", 1, true) ~= nil
-        return above_val or cur_val
-    end
     local function adjust_table() -- Add error to table
-        if check_spellwarn_comment() then return end
+        if M.check_spellwarn_comment(bufnr, vim.fn.line(".")) then return end
         local word = vim.fn.expand("<cword>")
         table.insert(errors, {
             col  = location[3],
@@ -63,4 +79,62 @@ function M.get_spelling_errors(bufnr)
     return errors
 end
 
+function M.get_spelling_errors_iter(bufnr, start_row, start_col, end_row, end_col)
+    if start_row == nil then start_row = 0 end
+    if start_col == nil then start_col = 0 end
+    if end_row == nil then end_row = #(vim.api.nvim_buf_get_lines(bufnr, 1, -1, false)) + 1 end
+    if end_col == nil then end_col = string.len(vim.api.nvim_buf_get_lines(bufnr, end_row - 1, end_row, false)[1]) end
+    local lines = vim.api.nvim_buf_get_lines(bufnr, start_row, end_row + 1, false)
+    lines[#lines] = string.sub(lines[#lines], 1, end_col + 1)
+    lines[1] = string.sub(lines[1], start_col + 1)
+    local errors = {}
+    for n, line in ipairs(lines) do
+        local errs = vim.spell.check(line)
+        for _, err in ipairs(errs) do
+            local i = start_row + n
+            local offset = (n == 1 and start_col) or 0
+            local key = i .. (err[3] + offset) -- By inserting based on location, we avoid duplicates
+            if not M.check_spellwarn_comment(bufnr, i) then
+                errors[key] = {
+                    lnum = i,
+                    col = err[3] + offset,
+                    word = err[1],
+                    type = "spell" .. err[2],
+                }
+            end
+        end
+    end
+    return errors
+end
+
+function M.get_spelling_errors_ts(bufnr)
+    local errors = {}
+    local ts_enabled = pcall(require, "nvim-treesitter")
+    local buf_highlighter = ts_enabled and vim.treesitter.highlighter.active[bufnr]
+
+    if not buf_highlighter then return M.get_spelling_errors_iter(bufnr) end
+    buf_highlighter.tree:for_each_tree(function(tstree, tree)
+        ---@diagnostic disable: invisible
+        if not tstree then return end
+        local root = tstree:root()
+
+        local q = buf_highlighter:get_query(tree:lang())
+
+        -- Some injected languages may not have highlight queries.
+        if not q:query() then return end
+
+        for capture, node in q:query():iter_captures(root, bufnr, 0, -1) do
+            local c = q._query.captures[capture] -- Name of the capture in the query
+            if c == "spell" then
+                local start_row, start_col, end_row, end_col = node:range()
+                for k, v in pairs(M.get_spelling_errors_iter(bufnr, start_row, start_col, end_row, end_col)) do
+                    errors[k] = v
+                end
+            end
+        end
+        ---@diagnostic enable: invisible
+    end)
+    return errors
+end
+
 return M