neovim

commit 302d3cfb96d7f0c856262e1a4252d058e3300c8b
parent b302da9ad220a7699d4b0ebf642529d142a0b9cf
Author: Mathias Fußenegger <mfussenegger@users.noreply.github.com>
Date:   Sat, 10 Jun 2023 20:33:23 +0200

feat(lua): use callable table as iterator in vim.iter (#23957)

A table passed to `vim.iter` can be a class instance with a `__call`
implementation for the iterator protocol.
Diffstat:
M
runtime/doc/lua.txt
 | 
8
++++++++
M
runtime/lua/vim/iter.lua
 | 
19
+++++++++++++++++++
M
test/functional/lua/iter_spec.lua
 | 
9
+++++++++

3 files changed, 36 insertions(+), 0 deletions(-)
diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
@@ -2999,6 +2999,8 @@ pipeline depend on the type passed to this function:
 • Non-list tables pass both the key and value of each element
 • Function iterators pass all of the values returned by their respective
   function
+• Tables with a metatable implementing __call are treated as function
+  iterators
 
 Examples: >lua
 
@@ -3032,6 +3034,12 @@ Examples: >lua
    end)
    -- true
 
+   local rb = vim.ringbuf(3)
+   rb:push("a")
+   rb:push("b")
+   vim.iter(rb):totable()
+   -- { "a", "b" }
+
 <
 
 In addition to the |vim.iter()| function, the |vim.iter| module provides
diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua
@@ -14,6 +14,8 @@
 --- - Non-list tables pass both the key and value of each element
 --- - Function iterators pass all of the values returned by their respective
 ---   function
+--- - Tables with a metatable implementing __call are treated as function
+---   iterators
 ---
 --- Examples:
 --- <pre>lua
@@ -46,6 +48,12 @@
 ---     return k == 'z'
 ---   end)
 ---   -- true
+---
+---   local rb = vim.ringbuf(3)
+---   rb:push("a")
+---   rb:push("b")
+---   vim.iter(rb):totable()
+---   -- { "a", "b" }
 --- </pre>
 ---
 --- In addition to the |vim.iter()| function, the |vim.iter| module provides
@@ -889,6 +897,17 @@ end
 function Iter.new(src, ...)
   local it = {}
   if type(src) == 'table' then
+    local mt = getmetatable(src)
+    if mt and type(mt.__call) == 'function' then
+      ---@private
+      function it.next()
+        return src()
+      end
+
+      setmetatable(it, Iter)
+      return it
+    end
+
     local t = {}
 
     -- Check if source table can be treated like a list (indices are consecutive integers
diff --git a/test/functional/lua/iter_spec.lua b/test/functional/lua/iter_spec.lua
@@ -4,6 +4,15 @@ local matches = helpers.matches
 local pcall_err = helpers.pcall_err
 
 describe('vim.iter', function()
+  it('new() on iterable class instance', function()
+    local rb = vim.ringbuf(3)
+    rb:push("a")
+    rb:push("b")
+
+    local it = vim.iter(rb)
+    eq({"a", "b"}, it:totable())
+  end)
+
   it('filter()', function()
     local function odd(v)
       return v % 2 ~= 0