feat(#2789): add optional function expand_until to api.tree.expand_all and api.node.expand (#3166)

* feat: Allow to expand nodes until certain condition is met

* Fix warnings

* Restore original position of edit function

* Rename field to match the api method name

* Rename ApiTreeExpandAllOpts to ApiTreeExpandOpts

* Remove toggle_descend_until

* Remove redundant empty line

* Update :help for changed methods

* Fix partial expansion of grouped nodes

* Fix lint error

* Fix linting error

* Fix incorrect open/close indicator state

* Update docs

* Rename descend_until option to expand_until

* Always check directory expansion limit

* Fix linter errors

* Ignore unused param warning

* Apply suggestions from code review

* simplify MAX_FOLDER_DISCOVERY warning

* fix bad comment whitespace

---------

Co-authored-by: ghostbuster91 <ghostbuster91@users.noreply.github.com>
Co-authored-by: Alexander Courtis <alex@courtis.org>

doc/nvim-tree-lua.txt

@@ -1841,11 +1841,17 @@ tree.collapse_all({opts})                  *nvim-tree-api.tree.collapse_all()*
     Options: ~
       • {keep_buffers} (boolean) do not collapse nodes with open buffers.
 
-tree.expand_all({node})                      *nvim-tree-api.tree.expand_all()*
+tree.expand_all({node}, {opts})              *nvim-tree-api.tree.expand_all()*
     Recursively expand all nodes under the tree root or specified folder.
 
     Parameters: ~
       • {node} (Node|nil) folder
+      • {opts} (ApiTreeExpandOpts) optional parameters
+
+    Options: ~
+      • {expand_until} ((fun(expansion_count: integer, node: Node?): boolean)?)
+            Return true if {node} should be expanded.
+            {expansion_count} is the total number of folders expanded.
 
                                   *nvim-tree-api.tree.toggle_enable_filters()*
 tree.toggle_enable_filters()
@@ -2279,12 +2285,18 @@ node.buffer.wipe({node}, {opts})            *nvim-tree-api.node.buffer.wipe()*
     Options: ~
       • {force} (boolean) wipe even if buffer is modified, default false
 
-node.expand({node})                              *nvim-tree-api.node.expand()*
+node.expand({node}, {opts})                      *nvim-tree-api.node.expand()*
     Recursively expand all nodes under a directory or a file's parent
     directory.
 
     Parameters: ~
       • {node} (Node|nil) file or folder
+      • {opts} (ApiTreeExpandOpts) optional parameters
+
+    Options: ~
+      • {expand_until} ((fun(expansion_count: integer, node: Node?): boolean)?)
+            Return true if {node} should be expanded.
+            {expansion_count} is the total number of folders expanded.
 
 node.collapse({node}, {opts})                  *nvim-tree-api.node.collapse()*
     Collapse the tree under a directory or a file's parent directory.

lua/nvim-tree/actions/tree/modifiers/expand.lua

@@ -27,20 +27,66 @@ local function expand(node)
   end
 end
 
----@param expansion_count integer
+---@param should_descend fun(expansion_count: integer, node: Node): boolean
+---@return fun(expansion_count: integer, node: Node): boolean
+local function limit_folder_discovery(should_descend)
+  return function(expansion_count, node)
+    local should_halt = expansion_count >= M.MAX_FOLDER_DISCOVERY
+    if should_halt then
+      notify.warn("expansion iteration was halted after " .. M.MAX_FOLDER_DISCOVERY .. " discovered folders")
+      return false
+    end
+
+    return should_descend(expansion_count, node)
+  end
+end
+
+---@param _ integer expansion_count
 ---@param node Node
 ---@return boolean
-local function should_expand(expansion_count, node)
+local function descend_until_empty(_, node)
+
   local dir = node:as(DirectoryNode)
   if not dir then
     return false
   end
-  local should_halt = expansion_count >= M.MAX_FOLDER_DISCOVERY
+
   local should_exclude = M.EXCLUDE[dir.name]
-  return not should_halt and not dir.open and not should_exclude
+  return not should_exclude
 end
 
-local function gen_iterator()
+---@param expansion_count integer
+---@param node Node
+---@param should_descend fun(expansion_count: integer, node: Node): boolean
+---@return boolean
+local function should_expand(expansion_count, node, should_descend)
+  local dir = node:as(DirectoryNode)
+  if not dir then
+    return false
+  end
+
+  if not dir.open and should_descend(expansion_count, node) then
+    if #node.nodes == 0 then
+      core.get_explorer():expand(dir) -- populate node.group_next
+    end
+
+    if dir.group_next then
+      local expand_next = should_expand(expansion_count, dir.group_next, should_descend)
+      if expand_next then
+        dir.open = true
+      end
+      return expand_next
+    else
+      return true
+    end
+  end
+  return false
+end
+
+
+---@param should_descend fun(expansion_count: integer, node: Node): boolean
+---@return fun(node): any
+local function gen_iterator(should_descend)
   local expansion_count = 0
 
   return function(parent)
@@ -52,7 +98,7 @@ local function gen_iterator()
     Iterator.builder(parent.nodes)
       :hidden()
       :applier(function(node)
-        if should_expand(expansion_count, node) then
+        if should_expand(expansion_count, node, should_descend) then
           expansion_count = expansion_count + 1
           node = node:as(DirectoryNode)
           if node then
@@ -61,25 +107,32 @@ local function gen_iterator()
         end
       end)
       :recursor(function(node)
-        return expansion_count < M.MAX_FOLDER_DISCOVERY and (node.group_next and { node.group_next } or (node.open and node.nodes))
+        if not should_descend(expansion_count, node) then
+          return nil
+        end
+
+        if node.group_next then
+          return { node.group_next }
+        end
+
+        if node.open and node.nodes then
+          return node.nodes
+        end
+
+        return nil
       end)
       :iterate()
-
-    if expansion_count >= M.MAX_FOLDER_DISCOVERY then
-      return true
-    end
   end
 end
 
 ---@param node Node?
-local function expand_node(node)
+---@param expand_opts ApiTreeExpandOpts?
+local function expand_node(node, expand_opts)
   if not node then
     return
   end
-
-  if gen_iterator()(node) then
-    notify.warn("expansion iteration was halted after " .. M.MAX_FOLDER_DISCOVERY .. " discovered folders")
-  end
+  local descend_until = limit_folder_discovery((expand_opts and expand_opts.expand_until) or descend_until_empty)
+  gen_iterator(descend_until)(node)
 
   local explorer = core.get_explorer()
   if explorer then
@@ -89,18 +142,20 @@ end
 
 ---Expand the directory node or the root
 ---@param node Node
-function M.all(node)
-  expand_node(node and node:as(DirectoryNode) or core.get_explorer())
+---@param expand_opts ApiTreeExpandOpts?
+function M.all(node, expand_opts)
+  expand_node(node and node:as(DirectoryNode) or core.get_explorer(), expand_opts)
 end
 
 ---Expand the directory node or parent node
 ---@param node Node
-function M.node(node)
+---@param expand_opts ApiTreeExpandOpts?
+function M.node(node, expand_opts)
   if not node then
     return
   end
 
-  expand_node(node:is(FileNode) and node.parent or node:as(DirectoryNode))
+  expand_node(node:is(FileNode) and node.parent or node:as(DirectoryNode), expand_opts)
 end
 
 function M.setup(opts)

lua/nvim-tree/api.lua

@@ -187,6 +187,10 @@ Api.tree.search_node = wrap(actions.finders.search_node.fn)
 ---@field keep_buffers boolean|nil default false
 
 Api.tree.collapse_all = wrap(actions.tree.modifiers.collapse.all)
+
+---@class ApiTreeExpandOpts
+---@field expand_until (fun(expansion_count: integer, node: Node): boolean)|nil
+
 Api.tree.expand_all = wrap_node(actions.tree.modifiers.expand.all)
 Api.tree.toggle_enable_filters = wrap_explorer_member("filters", "toggle")
 Api.tree.toggle_gitignore_filter = wrap_explorer_member_args("filters", "toggle", "git_ignored")