feat(lua): add vim.spell (#16620)

lewis6991 · web-flow · commit e11a44aa224a · 2021-12-25T12:36:56.000-07:00
diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
@@ -707,6 +707,38 @@ vim.mpack.encode({obj})                                 *vim.mpack.encode*
 vim.mpack.decode({str})                                 *vim.mpack.decode*
         Decodes (or "unpacks") the msgpack-encoded {str} to a Lua object.
 
+------------------------------------------------------------------------------
+VIM.SPELL                                                           *lua-spell*
+
+vim.spell.check({str})                                      *vim.spell.check()*
+        Check {str} for spelling errors. Similar to the Vimscript function
+        |spellbadword()|.
+
+        Note: The behaviour of this function is dependent on: 'spelllang',
+        'spellfile', 'spellcapcheck' and 'spelloptions' which can all be local
+        to the buffer. Consider calling this with |nvim_buf_call()|.
+
+        Example: >
+            vim.spell.check("the quik brown fox")
+            -->
+            {
+                {'quik', 'bad', 4}
+            }
+<
+
+        Parameters: ~
+            {str}    String to spell check.
+
+        Return: ~
+          List of tuples with three items:
+            - The badly spelled word.
+            - The type of the spelling error:
+                "bad"   spelling mistake
+                "rare"  rare word
+                "local" word only valid in another region
+                "caps"  word should start with Capital
+            - The position in {str} where the word begins.
+
 ------------------------------------------------------------------------------
 VIM                                                     *lua-builtin*
 
diff --git a/src/nvim/lua/spell.c b/src/nvim/lua/spell.c
@@ -0,0 +1,99 @@
+
+#include <lua.h>
+#include <lauxlib.h>
+
+#include "nvim/spell.h"
+#include "nvim/vim.h"
+#include "nvim/lua/spell.h"
+
+#ifdef INCLUDE_GENERATED_DECLARATIONS
+# include "lua/spell.c.generated.h"
+#endif
+
+int nlua_spell_check(lua_State *lstate)
+{
+  if (lua_gettop(lstate) < 1) {
+    return luaL_error(lstate, "Expected 1 argument");
+  }
+
+  if (lua_type(lstate, 1) != LUA_TSTRING) {
+    luaL_argerror(lstate, 1, "expected string");
+  }
+
+  const char *str = lua_tolstring(lstate, 1, NULL);
+
+  // spell.c requires that 'spell' is enabled, so we need to temporarily enable
+  // it before we can call spell functions.
+  const int wo_spell_save = curwin->w_p_spell;
+
+  if (!curwin->w_p_spell) {
+    did_set_spelllang(curwin);
+    curwin->w_p_spell = true;
+  }
+
+  // Check 'spelllang'
+  if (*curwin->w_s->b_p_spl == NUL) {
+    emsg(_(e_no_spell));
+    curwin->w_p_spell = wo_spell_save;
+    return 0;
+  }
+
+  hlf_T attr = HLF_COUNT;
+  size_t len = 0;
+  size_t pos = 0;
+  int capcol = -1;
+  int no_res = 0;
+  const char * result;
+
+  lua_createtable(lstate, 0, 0);
+
+  while (*str != NUL) {
+    attr = HLF_COUNT;
+    len = spell_check(curwin, (char_u *)str, &attr, &capcol, false);
+    assert(len <= INT_MAX);
+
+    if (attr != HLF_COUNT) {
+      lua_createtable(lstate, 3, 0);
+
+      lua_pushlstring(lstate, str, len);
+      lua_rawseti(lstate, -2, 1);
+
+      result = attr == HLF_SPB ? "bad"   :
+               attr == HLF_SPR ? "rare"  :
+               attr == HLF_SPL ? "local" :
+               attr == HLF_SPC ? "caps"  :
+               NULL;
+
+      assert(result != NULL);
+
+      lua_pushstring(lstate, result);
+      lua_rawseti(lstate, -2, 2);
+
+      // +1 for 1-indexing
+      lua_pushinteger(lstate, (long)pos + 1);
+      lua_rawseti(lstate, -2, 3);
+
+      lua_rawseti(lstate, -2, ++no_res);
+    }
+
+    str += len;
+    pos += len;
+    capcol -= (int)len;
+  }
+
+  // Restore 'spell'
+  curwin->w_p_spell = wo_spell_save;
+  return 1;
+}
+
+static const luaL_Reg spell_functions[] = {
+  { "check", nlua_spell_check },
+  { NULL   , NULL }
+};
+
+int luaopen_spell(lua_State *L)
+{
+  lua_newtable(L);
+  luaL_register(L, NULL, spell_functions);
+  return 1;
+}
diff --git a/src/nvim/lua/spell.h b/src/nvim/lua/spell.h
@@ -0,0 +1,12 @@
+#ifndef NVIM_LUA_SPELL_H
+#define NVIM_LUA_SPELL_H
+
+#include <lauxlib.h>
+#include <lua.h>
+#include <lualib.h>
+
+#ifdef INCLUDE_GENERATED_DECLARATIONS
+# include "lua/spell.h.generated.h"
+#endif
+
+#endif  // NVIM_LUA_SPELL_H
diff --git a/src/nvim/lua/stdlib.c b/src/nvim/lua/stdlib.c
@@ -30,6 +30,7 @@
 #include "nvim/lua/stdlib.h"
 #include "nvim/lua/treesitter.h"
 #include "nvim/lua/xdiff.h"
+#include "nvim/lua/spell.h"
 #include "nvim/macros.h"
 #include "nvim/map.h"
 #include "nvim/memline.h"
@@ -518,6 +519,10 @@ void nlua_state_add_stdlib(lua_State *const lstate)
   lua_pushcfunction(lstate, &nlua_xdl_diff);
   lua_setfield(lstate, -2, "diff");
 
+  // vim.spell
+  luaopen_spell(lstate);
+  lua_setfield(lstate, -2, "spell");
+
   lua_cjson_new(lstate);
   lua_setfield(lstate, -2, "json");
 }
diff --git a/test/functional/lua/spell_spec.lua b/test/functional/lua/spell_spec.lua
@@ -0,0 +1,53 @@
+local helpers = require('test.functional.helpers')(after_each)
+local clear = helpers.clear
+local exec_lua = helpers.exec_lua
+local eq = helpers.eq
+local pcall_err = helpers.pcall_err
+
+describe('vim.spell', function()
+  before_each(function()
+    clear()
+  end)
+
+  describe('.check', function()
+    local check = function(x, exp)
+        return eq(exp, exec_lua("return vim.spell.check(...)", x))
+    end
+
+    it('can handle nil', function()
+      eq([[Error executing lua: [string "<nvim>"]:0: bad argument #1 to 'check' (expected string)]],
+        pcall_err(exec_lua, [[vim.spell.check(nil)]]))
+    end)
+
+    it('can check spellings', function()
+      check('hello', {})
+
+      check(
+        'helloi',
+        {{"helloi", "bad", 1}}
+      )
+
+      check(
+        'hello therei',
+        {{"therei", "bad", 7}}
+      )
+
+      check(
+        'hello. there',
+        {{"there", "caps", 8}}
+      )
+
+      check(
+        'neovim cna chkc spellins. okay?',
+        {
+          {"neovim"  , "bad" ,  1},
+          {"cna"     , "bad" ,  8},
+          {"chkc"    , "bad" , 12},
+          {"spellins", "bad" , 17},
+          {"okay"    , "caps", 27}
+        }
+      )
+    end)
+
+  end)
+end)
