feat: Allow to expand until certain condition is met

doc/nvim-tree-lua.txt

@@ -1808,10 +1808,35 @@ tree.collapse_all({keep_buffers})          *nvim-tree-api.tree.collapse_all()*
     Parameters: ~
       • {keep_buffers} (boolean) do not collapse nodes with open buffers.
 
-tree.expand_all()                            *nvim-tree-api.tree.expand_all()*
+tree.expand_all({opts})                      *nvim-tree-api.tree.expand_all()*
     Recursively expand all nodes in the tree.
     Folder: only the nodes underneath that folder.
 
+    Parameters: ~
+      • {opts} (table) optional parameters
+
+    Options: ~
+      • {expand_until} (fun(expansion_count: integer, node: Node): boolean)
+
+            A function returning boolean that specifies terminating condition
+            for node/tree expansion. If not provided the function that expands
+            recursively entire node/tree will be used.
+
+        Example: >
+          -- expand only 5 levels deep
+          local function my_expand_until(expansion_count, node)
+            print("my_expand_until " .. expansion_count .. " " .. tostring(node and node.absolute_path))
+            local should_halt = expansion_count >= 5
+            return not should_halt
+          end
+
+          local function my_expand_all()
+            api.tree.expand_all(nil, { expand_until = my_expand_until })
+          end
+
+          -- on_attach
+          vim.keymap.set("n", "Z", my_expand_all, opts("My Expand All"))
+<
                                   *nvim-tree-api.tree.toggle_enable_filters()*
 tree.toggle_enable_filters()
     Toggle |nvim-tree.filters.enable| all filters.

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

@@ -18,25 +18,38 @@ local function to_lookup_table(list)
 end
 
 ---@param node Node
-local function expand(node)
-  node = lib.get_last_group_node(node)
-  node.open = true
+local function populate_node(node)
+  -- noop if it is a file
+  if node.nodes == nil then
+    return
+  end
   if #node.nodes == 0 then
+    local cwd = node.link_to or node.absolute_path
+    local handle = vim.loop.fs_scandir(cwd)
+    if not handle then
+      return
+    end
     core.get_explorer():expand(node)
   end
 end
 
 ---@param expansion_count integer
 ---@param node Node
 ---@return boolean
-local function should_expand(expansion_count, node)
+local function expand_until_max_or_empty(expansion_count, node)
   local should_halt = expansion_count >= M.MAX_FOLDER_DISCOVERY
   local should_exclude = M.EXCLUDE[node.name]
   return not should_halt and node.nodes and not node.open and not should_exclude
 end
 
-local function gen_iterator()
+---@param expand_until fun(expansion_count: integer, node: Node): boolean
+local function gen_iterator(expand_until)
   local expansion_count = 0
+  local function expand(node)
+    populate_node(node)
+    node = lib.get_last_group_node(node)
+    node.open = true
+  end
 
   return function(parent)
     if parent.parent and parent.nodes and not parent.open then
@@ -47,13 +60,14 @@ local function gen_iterator()
     Iterator.builder(parent.nodes)
       :hidden()
       :applier(function(node)
-        if should_expand(expansion_count, node) then
+        if expand_until(expansion_count, node, populate_node) then
           expansion_count = expansion_count + 1
           expand(node)
         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))
+        local should_recurse = expand_until(expansion_count - 1, node, populate_node)
+        return expansion_count < M.MAX_FOLDER_DISCOVERY and should_recurse and node.nodes
       end)
       :iterate()
 
@@ -64,9 +78,11 @@ local function gen_iterator()
 end
 
 ---@param base_node table
-function M.fn(base_node)
+---@param expand_opts ApiTreeExpandAllOpts|nil
+function M.fn(base_node, expand_opts)
+  local expand_until = (expand_opts and expand_opts.expand_until) or expand_until_max_or_empty
   local node = base_node.nodes and base_node or core.get_explorer()
-  if gen_iterator()(node) then
+  if gen_iterator(expand_until)(node) then
     notify.warn("expansion iteration was halted after " .. M.MAX_FOLDER_DISCOVERY .. " discovered folders")
   end
   renderer.draw()

lua/nvim-tree/api.lua

@@ -144,7 +144,12 @@ Api.tree.get_nodes = wrap(lib.get_nodes)
 Api.tree.find_file = wrap(actions.tree.find_file.fn)
 Api.tree.search_node = wrap(actions.finders.search_node.fn)
 Api.tree.collapse_all = wrap(actions.tree.modifiers.collapse_all.fn)
+
+---@class ApiTreeExpandAllOpts
+---@field expand_until (fun(expansion_count: integer, node: Node): boolean)|nil
+
 Api.tree.expand_all = wrap_node(actions.tree.modifiers.expand_all.fn)
+
 Api.tree.toggle_enable_filters = wrap(actions.tree.modifiers.toggles.enable)
 Api.tree.toggle_gitignore_filter = wrap(actions.tree.modifiers.toggles.git_ignored)
 Api.tree.toggle_git_clean_filter = wrap(actions.tree.modifiers.toggles.git_clean)

lua/nvim-tree/explorer/node.lua

@@ -36,7 +36,7 @@ end
 ---@param node Node
 ---@return boolean
 function M.has_one_child_folder(node)
-  return #node.nodes == 1 and node.nodes[1].nodes and vim.loop.fs_access(node.nodes[1].absolute_path, "R") or false
+  return node.nodes ~= nil and #node.nodes == 1 and node.nodes[1].nodes and vim.loop.fs_access(node.nodes[1].absolute_path, "R") or false
 end
 
 ---@param node Node