feat: better support for plugin actions with custom lua function

lua/which-key/keys.lua

@@ -80,11 +80,9 @@ function M.get_mappings(mode, prefix, buf)
     end
   end
 
-  add(M.get_tree(mode).tree:get(prefix))
-  add(M.get_tree(mode, buf).tree:get(prefix))
-
-  -- Run a plugin if needed
-  if ret.mapping and ret.mapping.plugin then require("which-key.plugins").invoke(ret) end
+  local plugin_context = { buf = buf, mode = mode }
+  add(M.get_tree(mode).tree:get(prefix, nil, plugin_context))
+  add(M.get_tree(mode, buf).tree:get(prefix, nil, plugin_context))
 
   -- Handle motions
   M.process_motions(ret, mode, prefix, buf)
@@ -107,20 +105,19 @@ function M.get_mappings(mode, prefix, buf)
   end
 
   -- Sort items, but not for plugins
-  if not (ret.mapping and ret.mapping.plugin) then
-    table.sort(tmp, function(a, b)
-      if a.group == b.group then
-        local ak = (a.key or ""):lower()
-        local bk = (b.key or ""):lower()
-        local aw = ak:match("[a-z]") and 1 or 0
-        local bw = bk:match("[a-z]") and 1 or 0
-        if aw == bw then return ak < bk end
-        return aw < bw
-      else
-        return (a.group and 1 or 0) < (b.group and 1 or 0)
-      end
-    end)
-  end
+  table.sort(tmp, function(a, b)
+    if a.order and b.order then return a.order < b.order end
+    if a.group == b.group then
+      local ak = (a.key or ""):lower()
+      local bk = (b.key or ""):lower()
+      local aw = ak:match("[a-z]") and 1 or 0
+      local bw = bk:match("[a-z]") and 1 or 0
+      if aw == bw then return ak < bk end
+      return aw < bw
+    else
+      return (a.group and 1 or 0) < (b.group and 1 or 0)
+    end
+  end)
   ret.mappings = tmp
 
   return ret

lua/which-key/plugins/init.lua

@@ -34,25 +34,21 @@ function M._setup(plugin, opts)
   if plugin.setup then plugin.setup(require("which-key"), opts, Config.options) end
 end
 
----@param results MappingGroup
-function M.invoke(results)
-  local plugin = M.plugins[results.mapping.plugin]
-  local prefix = results.mapping.prefix
-  local items = plugin.run(prefix, results.mode, results.buf)
-
-  for _, item in pairs(items) do
+---@param mapping Mapping
+function M.invoke(mapping, context)
+  local plugin = M.plugins[mapping.plugin]
+  local prefix = mapping.prefix
+  local items = plugin.run(prefix, context.mode, context.buf)
+
+  local ret = {}
+  for i, item in ipairs(items) do
     ---@type VisualMapping
-    local mapping
-    mapping = {
-      key = item.key,
-      label = item.label,
-      keys = Util.parse_keys(prefix .. item.key),
-      prefix = prefix,
-      value = item.value,
-      highlights = item.highlights,
-    }
-    table.insert(results.mappings, mapping)
+    item.order = i
+    item.keys = Util.parse_keys(prefix .. item.key)
+    item.prefix = prefix .. item.key
+    table.insert(ret, item)
   end
+  return ret
 end
 
 return M

lua/which-key/tree.lua

@@ -21,13 +21,18 @@ end
 ---@param prefix string
 ---@param index number defaults to last. If < 0, then offset from last
 ---@return Node
-function Tree:get(prefix, index)
+function Tree:get(prefix, index, plugin_context)
   prefix = Util.parse_keys(prefix).term
   local node = self.root
   index = index or #prefix
   if index < 0 then index = #prefix + index end
   for i = 1, index, 1 do
     node = node.children[prefix[i]]
+    if node and plugin_context and node.mapping and node.mapping.plugin then
+      local children = require("which-key.plugins").invoke(node.mapping, plugin_context)
+      node.children = {}
+      for _, child in pairs(children) do self:add(child) end
+    end
     if not node then return nil end
   end
   return node

lua/which-key/types.lua

@@ -34,6 +34,7 @@ local MappingOptions
 ---@field opts MappingOptions
 ---@field keys KeyCodes
 ---@field plugin string
+---@field fn fun()
 local Mapping
 
 ---@class MappingTree

lua/which-key/view.lua

@@ -193,7 +193,11 @@ function M.on_keys(opts)
     --- Check for an exact match. Feedkeys with remap
     if results.mapping and not results.mapping.group and #results.mappings == 0 then
       M.hide()
-      M.execute(M.keys, M.mode, buf)
+      if results.mapping.fn then
+        results.mapping.fn()
+      else
+        M.execute(M.keys, M.mode, buf)
+      end
       return
     end