feat(lua): allow vim.on_key() callback to consume the key (#30939)

1 files changed, 65 insertions, 0 deletions

diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index 13e146a9da..55ca489a9a 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -28,6 +28,7 @@ local rmdir = n.rmdir
 local write_file = t.write_file
 local poke_eventloop = n.poke_eventloop
 local assert_alive = n.assert_alive
+local expect = n.expect
 
 describe('lua stdlib', function()
   before_each(clear)
@@ -3416,6 +3417,70 @@ describe('lua stdlib', function()
                                                                     |
       ]])
     end)
+
+    it('can discard input', function()
+      clear()
+      -- discard every other normal 'x' command
+      exec_lua [[
+        n_key = 0
+
+        vim.on_key(function(buf, typed_buf)
+          if typed_buf == 'x' then
+            n_key = n_key + 1
+          end
+          return (n_key % 2 == 0) and "" or nil
+        end)
+      ]]
+
+      api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
+
+      feed('x')
+      expect('4321')
+      feed('x')
+      expect('4321')
+      feed('x')
+      expect('321')
+      feed('x')
+      expect('321')
+    end)
+
+    it('callback invalid return', function()
+      clear()
+      -- second key produces an error which removes the callback
+      exec_lua [[
+        n_call = 0
+        vim.on_key(function(buf, typed_buf)
+          if typed_buf == 'x' then
+            n_call = n_call + 1
+          end
+          return n_call >= 2 and '!' or nil
+        end)
+      ]]
+
+      api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
+
+      local function cleanup_msg(msg)
+        return (remove_trace(msg):gsub('^Error.*\n *Messages: ', ''))
+      end
+
+      feed('x')
+      eq(1, exec_lua [[ return n_call ]])
+
+      eq(1, exec_lua [[ return vim.on_key(nil, nil) ]])
+
+      eq('', cleanup_msg(eval('v:errmsg')))
+      feed('x')
+      eq(2, exec_lua [[ return n_call ]])
+      eq('return string must be empty', cleanup_msg(eval('v:errmsg')))
+      command('let v:errmsg = ""')
+
+      eq(0, exec_lua [[ return vim.on_key(nil, nil) ]])
+
+      feed('x')
+      eq(2, exec_lua [[ return n_call ]])
+      expect('21')
+      eq('', cleanup_msg(eval('v:errmsg')))
+    end)
   end)
 
   describe('vim.wait', function()