diff options
| -rw-r--r-- | runtime/doc/lua.txt | 22 | ||||
| -rw-r--r-- | runtime/lua/vim/iter.lua | 82 | ||||
| -rw-r--r-- | test/functional/lua/iter_spec.lua | 32 |
3 files changed, 132 insertions, 4 deletions
diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt index 1bae1a43d4..b558a3fc8d 100644 --- a/runtime/doc/lua.txt +++ b/runtime/doc/lua.txt @@ -3431,6 +3431,28 @@ Iter:find({f}) *Iter:find()* Return: ~ any +Iter:flatten({depth}) *Iter:flatten()* + Flattens a |list-iterator|, un-nesting nested values up to the given + {depth}. Errors if it attempts to flatten a dict-like value. + + Examples: >lua + vim.iter({ 1, { 2 }, { { 3 } } }):flatten():totable() + -- { 1, 2, { 3 } } + + vim.iter({1, { { a = 2 } }, { 3 } }):flatten():totable() + -- { 1, { a = 2 }, 3 } + + vim.iter({ 1, { { a = 2 } }, { 3 } }):flatten(math.huge):totable() + -- error: attempt to flatten a dict-like table +< + + Parameters: ~ + • {depth} (number|nil) Depth to which |list-iterator| should be + flattened (defaults to 1) + + Return: ~ + Iter + Iter:fold({init}, {f}) *Iter:fold()* Folds ("reduces") an iterator into a single value. diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua index b658dde099..c6feeea3dc 100644 --- a/runtime/lua/vim/iter.lua +++ b/runtime/lua/vim/iter.lua @@ -112,6 +112,35 @@ local function sanitize(t) return t end +--- Flattens a single list-like table. Errors if it attempts to flatten a +--- dict-like table +---@param v table table which should be flattened +---@param max_depth number depth to which the table should be flattened +---@param depth number current iteration depth +---@param result table output table that contains flattened result +---@return table|nil flattened table if it can be flattened, otherwise nil +local function flatten(v, max_depth, depth, result) + if depth < max_depth and type(v) == 'table' then + local i = 0 + for _ in pairs(v) do + i = i + 1 + + if v[i] == nil then + -- short-circuit: this is not a list like table + return nil + end + + if flatten(v[i], max_depth, depth + 1, result) == nil then + return nil + end + end + else + result[#result + 1] = v + end + + return result +end + --- Determine if the current iterator stage should continue. --- --- If any arguments are passed to this function, then return those arguments @@ -179,6 +208,54 @@ function ListIter.filter(self, f) return self end +--- Flattens a |list-iterator|, un-nesting nested values up to the given {depth}. +--- Errors if it attempts to flatten a dict-like value. +--- +--- Examples: +--- +--- ```lua +--- vim.iter({ 1, { 2 }, { { 3 } } }):flatten():totable() +--- -- { 1, 2, { 3 } } +--- +--- vim.iter({1, { { a = 2 } }, { 3 } }):flatten():totable() +--- -- { 1, { a = 2 }, 3 } +--- +--- vim.iter({ 1, { { a = 2 } }, { 3 } }):flatten(math.huge):totable() +--- -- error: attempt to flatten a dict-like table +--- ``` +--- +---@param depth? number Depth to which |list-iterator| should be flattened +--- (defaults to 1) +---@return Iter +function Iter.flatten(self, depth) -- luacheck: no unused args + error('flatten() requires a list-like table') +end + +---@private +function ListIter.flatten(self, depth) + depth = depth or 1 + local inc = self._head < self._tail and 1 or -1 + local target = {} + + for i = self._head, self._tail - inc, inc do + local flattened = flatten(self._table[i], depth, 0, {}) + + -- exit early if we try to flatten a dict-like table + if flattened == nil then + error('flatten() requires a list-like table') + end + + for _, v in pairs(flattened) do + target[#target + 1] = v + end + end + + self._head = 1 + self._tail = #target + 1 + self._table = target + return self +end + --- Maps the items of an iterator pipeline to the values returned by `f`. --- --- If the map function returns nil, the value is filtered from the iterator. @@ -461,9 +538,8 @@ end --- ``` --- ---@return Iter -function Iter.rev(self) +function Iter.rev(self) -- luacheck: no unused args error('rev() requires a list-like table') - return self end ---@private @@ -733,7 +809,6 @@ end ---@diagnostic disable-next-line: unused-local function Iter.skipback(self, n) -- luacheck: no unused args error('skipback() requires a list-like table') - return self end ---@private @@ -800,7 +875,6 @@ end ---@diagnostic disable-next-line: unused-local function Iter.slice(self, first, last) -- luacheck: no unused args error('slice() requires a list-like table') - return self end ---@private diff --git a/test/functional/lua/iter_spec.lua b/test/functional/lua/iter_spec.lua index fdf573669a..8d6cf1264b 100644 --- a/test/functional/lua/iter_spec.lua +++ b/test/functional/lua/iter_spec.lua @@ -462,6 +462,38 @@ describe('vim.iter', function() ) end) + it('flatten()', function() + local t = { { 1, { 2 } }, { { { { 3 } } }, { 4 } }, { 5 } } + + eq(t, vim.iter(t):flatten(-1):totable()) + eq(t, vim.iter(t):flatten(0):totable()) + eq({ 1, { 2 }, { { { 3 } } }, { 4 }, 5 }, vim.iter(t):flatten():totable()) + eq({ 1, 2, { { 3 } }, 4, 5 }, vim.iter(t):flatten(2):totable()) + eq({ 1, 2, { 3 }, 4, 5 }, vim.iter(t):flatten(3):totable()) + eq({ 1, 2, 3, 4, 5 }, vim.iter(t):flatten(4):totable()) + + local m = { a = 1, b = { 2, 3 }, d = { 4 } } + local it = vim.iter(m) + + local flat_err = 'flatten%(%) requires a list%-like table' + matches(flat_err, pcall_err(it.flatten, it)) + + -- cases from the documentation + local simple_example = { 1, { 2 }, { { 3 } } } + eq({ 1, 2, { 3 } }, vim.iter(simple_example):flatten():totable()) + + local not_list_like = vim.iter({ [2] = 2 }) + matches(flat_err, pcall_err(not_list_like.flatten, not_list_like)) + + local also_not_list_like = vim.iter({ nil, 2 }) + matches(flat_err, pcall_err(not_list_like.flatten, also_not_list_like)) + + local nested_non_lists = vim.iter({ 1, { { a = 2 } }, { { nil } }, { 3 } }) + eq({ 1, { a = 2 }, { nil }, 3 }, nested_non_lists:flatten():totable()) + -- only error if we're going deep enough to flatten a dict-like table + matches(flat_err, pcall_err(nested_non_lists.flatten, nested_non_lists, math.huge)) + end) + it('handles map-like tables', function() local it = vim.iter({ a = 1, b = 2, c = 3 }):map(function(k, v) if v % 2 ~= 0 then |