tidy expand all

Author: alex-courtis

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

@@ -17,11 +17,9 @@ local function to_lookup_table(list)
   return table
 end
 
----@param node Node
+---@param node DirectoryNode
 local function expand(node)
-  if node:is(DirectoryNode) then
-    node = node:last_group_node()
-  end
+  node = node:last_group_node()
   node.open = true
   if #node.nodes == 0 then
     core.get_explorer():expand(node)
@@ -40,6 +38,7 @@ end
 local function gen_iterator()
   local expansion_count = 0
 
+  ---@param parent DirectoryNode
   return function(parent)
     if parent.parent and parent.nodes and not parent.open then
       expansion_count = expansion_count + 1
@@ -48,12 +47,14 @@ local function gen_iterator()
 
     Iterator.builder(parent.nodes)
       :hidden()
+      ---@param node DirectoryNode
       :applier(function(node)
         if should_expand(expansion_count, node) then
           expansion_count = expansion_count + 1
           expand(node)
         end
       end)
+      ---@param node DirectoryNode
       :recursor(function(node)
         return expansion_count < M.MAX_FOLDER_DISCOVERY and (node.group_next and { node.group_next } or (node.open and node.nodes))
       end)
@@ -65,11 +66,16 @@ local function gen_iterator()
   end
 end
 
+---Expand the directory node or the root
 ---@param node Node
 function M.fn(node)
   local explorer = core.get_explorer()
-  node = node.nodes and node or explorer
-  if gen_iterator()(node) then
+  local parent = node:as(DirectoryNode) or explorer
+  if not parent then
+    return
+  end
+
+  if gen_iterator()(parent) then
     notify.warn("expansion iteration was halted after " .. M.MAX_FOLDER_DISCOVERY .. " discovered folders")
   end
   if explorer then

lua/nvim-tree/node/directory.lua

@@ -133,7 +133,7 @@ function DirectoryNode:refresh()
 end
 
 -- If node is grouped, return the last node in the group. Otherwise, return the given node.
----@return Node
+---@return DirectoryNode
 function DirectoryNode:last_group_node()
   return self.group_next and self.group_next:last_group_node() or self
 end