feat(lua): function behavior for tbl_extend, tbl_deep_extend #33819

runtime/doc/lua.txt

@@ -2259,11 +2259,15 @@ vim.tbl_deep_extend({behavior}, {...})                 *vim.tbl_deep_extend()*
     overwritten instead of merged).
 
     Parameters: ~
-      • {behavior}  (`'error'|'keep'|'force'`) Decides what to do if a key is
+      • {behavior}  (`'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any`)
-                    found in more than one map:
+                    Decides what to do if a key is found in more than one map:
                     • "error": raise an error
                     • "keep": use value from the leftmost map
                     • "force": use value from the rightmost map
+                    • If a function, it receives the current key, value, and
+                      the previous value in the currently merged table (if
+                      present) and should return the value for the given key
+                      in the merged table.
       • {...}       (`table`) Two or more tables
 
     Return: ~
@@ -2276,11 +2280,15 @@ vim.tbl_extend({behavior}, {...})                           *vim.tbl_extend()*
     Merges two or more tables.
 
     Parameters: ~
-      • {behavior}  (`'error'|'keep'|'force'`) Decides what to do if a key is
+      • {behavior}  (`'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any`)
-                    found in more than one map:
+                    Decides what to do if a key is found in more than one map:
                     • "error": raise an error
                     • "keep": use value from the leftmost map
                     • "force": use value from the rightmost map
+                    • If a function, it receives the current key, value, and
+                      the previous value in the currently merged table (if
+                      present) and should return the value for the given key
+                      in the merged table.
       • {...}       (`table`) Two or more tables
 
     Return: ~

runtime/doc/news.txt

@@ -152,6 +152,7 @@ LUA
 
 • Lua type annotations for `vim.uv`.
 • |vim.hl.range()| now allows multiple timed highlights.
+• |vim.tbl_extend()| and |vim.tbl_deep_extend()| now accept a function behavior argument.
 
 OPTIONS
 

runtime/lua/vim/shared.lua

@@ -366,7 +366,7 @@ local function can_merge(v)
 end
 
 --- Recursive worker for tbl_extend
---- @param behavior 'error'|'keep'|'force'
+--- @param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any): any
 --- @param deep_extend boolean
 --- @param ... table<any,any>
 local function tbl_extend_rec(behavior, deep_extend, ...)
@@ -381,6 +381,8 @@ local function tbl_extend_rec(behavior, deep_extend, ...)
       for k, v in pairs(tbl) do
         if deep_extend and can_merge(v) and can_merge(ret[k]) then
           ret[k] = tbl_extend_rec(behavior, true, ret[k], v)
+        elseif type(behavior) == 'function' then
+          ret[k] = behavior(k, v, ret[k])
         elseif behavior ~= 'force' and ret[k] ~= nil then
           if behavior == 'error' then
             error('key found in more than one map: ' .. k)
@@ -395,11 +397,16 @@ local function tbl_extend_rec(behavior, deep_extend, ...)
   return ret
 end
 
---- @param behavior 'error'|'keep'|'force'
+--- @param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any): any
 --- @param deep_extend boolean
 --- @param ... table<any,any>
 local function tbl_extend(behavior, deep_extend, ...)
-  if behavior ~= 'error' and behavior ~= 'keep' and behavior ~= 'force' then
+  if
+    behavior ~= 'error'
+    and behavior ~= 'keep'
+    and behavior ~= 'force'
+    and type(behavior) ~= 'function'
+  then
     error('invalid "behavior": ' .. tostring(behavior))
   end
 
@@ -420,10 +427,12 @@ end
 ---
 ---@see |extend()|
 ---
----@param behavior 'error'|'keep'|'force' Decides what to do if a key is found in more than one map:
+---@param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any Decides what to do if a key is found in more than one map:
 ---      - "error": raise an error
 ---      - "keep":  use value from the leftmost map
 ---      - "force": use value from the rightmost map
+---      - If a function, it receives the current key, value, and the previous value in the currently merged table (if present) and should
+---        return the value for the given key in the merged table.
 ---@param ... table Two or more tables
 ---@return table : Merged table
 function vim.tbl_extend(behavior, ...)
@@ -441,10 +450,12 @@ end
 ---
 ---@generic T1: table
 ---@generic T2: table
----@param behavior 'error'|'keep'|'force' Decides what to do if a key is found in more than one map:
+---@param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any Decides what to do if a key is found in more than one map:
 ---      - "error": raise an error
 ---      - "keep":  use value from the leftmost map
 ---      - "force": use value from the rightmost map
+---      - If a function, it receives the current key, value, and the previous value in the currently merged table (if present) and should
+---        return the value for the given key in the merged table.
 ---@param ... T2 Two or more tables
 ---@return T1|T2 (table) Merged table
 function vim.tbl_deep_extend(behavior, ...)

test/functional/lua/vim_spec.lua

@@ -1036,6 +1036,20 @@ describe('lua stdlib', function()
       return c.x.a == 1 and c.x.b == 2 and c.x.c == nil and count == 1
     ]]))
 
+    ok(exec_lua([[
+      local a = { a = 1, b = 2, c = 1 }
+      local b = { a = -1, b = 5, c = 3, d = 4 }
+      -- Return the maximum value for each key.
+      local c = vim.tbl_extend(function(k, v, prev_v)
+        if prev_v then
+          return v > prev_v and v or prev_v
+        else
+          return v
+        end
+      end, a, b)
+      return vim.deep_equal(c, { a = 1, b = 5, c = 3, d = 4 })
+    ]]))
+
     matches(
       'invalid "behavior": nil',
       pcall_err(
@@ -1177,6 +1191,20 @@ describe('lua stdlib', function()
       return vim.deep_equal(c, { sub = { 'b', 'c' } })
     ]]))
 
+    ok(exec_lua([[
+      local a = { a = 1, b = 2, c = { d = 1, e = -2} }
+      local b = { a = -1, b = 5, c = { d = 6 } }
+      -- Return the maximum value for each key.
+      local c = vim.tbl_deep_extend(function(k, v, prev_v)
+        if prev_v then
+          return v > prev_v and v or prev_v
+        else
+          return v
+        end
+      end, a, b)
+      return vim.deep_equal(c, { a = 1, b = 5, c = { d = 6, e = -2 } })
+    ]]))
+
     matches('invalid "behavior": nil', pcall_err(exec_lua, [[return vim.tbl_deep_extend()]]))
 
     matches(