fix(commands): generalize expand_all_nodes (#1661)

Author: pynappo

Diff for: lua/neo-tree/sources/common/commands.lua

@@ -1,5 +1,4 @@
 --This file should contain all commands meant to be used by mappings.
-
 local vim = vim
 local fs_actions = require("neo-tree.sources.filesystem.lib.fs_actions")
 local utils = require("neo-tree.utils")
@@ -113,17 +112,17 @@ end
 ---Expand all nodes
 ---@param state table The state of the source
 ---@param node table A node to expand
----@param prefetcher table an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
+---@param prefetcher table? an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
 M.expand_all_nodes = function(state, node, prefetcher)
-  log.debug("Expanding all nodes under " .. node:get_id())
-  if prefetcher == nil then
-    prefetcher = node_expander.default_prefetcher
-  end
+  local root_nodes = node and { node } or state.tree:get_nodes()
 
   renderer.position.set(state, nil)
 
   local task = function()
-    node_expander.expand_directory_recursively(state, node, prefetcher)
+    for _, root in pairs(root_nodes) do
+      log.debug("Expanding all nodes under " .. root:get_id())
+      node_expander.expand_directory_recursively(state, root, prefetcher)
+    end
   end
   async.run(task, function()
     log.debug("All nodes expanded - redrawing")
@@ -150,11 +149,8 @@ M.close_node = function(state, callback)
     target_node:collapse()
     renderer.redraw(state)
     renderer.focus_node(state, target_node:get_id())
-    if
-      state.explicitly_opened_directories
-      and state.explicitly_opened_directories[target_node:get_id()]
-    then
-      state.explicitly_opened_directories[target_node:get_id()] = false
+    if state.explicitly_opened_nodes and state.explicitly_opened_nodes[target_node:get_id()] then
+      state.explicitly_opened_nodes[target_node:get_id()] = false
     end
   end
 end
@@ -174,16 +170,13 @@ M.close_all_subnodes = function(state)
   renderer.collapse_all_nodes(tree, target_node:get_id())
   renderer.redraw(state)
   renderer.focus_node(state, target_node:get_id())
-  if
-    state.explicitly_opened_directories
-    and state.explicitly_opened_directories[target_node:get_id()]
-  then
-    state.explicitly_opened_directories[target_node:get_id()] = false
+  if state.explicitly_opened_nodes and state.explicitly_opened_nodes[target_node:get_id()] then
+    state.explicitly_opened_nodes[target_node:get_id()] = false
   end
 end
 
 M.close_all_nodes = function(state)
-  state.explicitly_opened_directories = {}
+  state.explicitly_opened_nodes = {}
   renderer.collapse_all_nodes(state.tree)
   renderer.redraw(state)
 end

Diff for: lua/neo-tree/sources/common/node_expander.lua

@@ -1,4 +1,5 @@
 local log = require("neo-tree.log")
+local utils = require("neo-tree.utils")
 
 local M = {}
 
@@ -15,15 +16,15 @@ local function expand_loaded(node, state, prefetcher)
     else
       if not current_node:is_expanded() then
         current_node:expand()
-        state.explicitly_opened_directories[current_node:get_id()] = true
+        state.explicitly_opened_nodes[current_node:get_id()] = true
       end
       local children = state.tree:get_nodes(current_node:get_id())
       log.debug("Expanding childrens of " .. current_node:get_id())
       for _, child in ipairs(children) do
-        if child.type == "directory" then
+        if utils.is_expandable(child) then
           rec(child, to_load)
         else
-          log.trace("Child: " .. child.name .. " is not a directory, skipping")
+          log.trace("Child: " .. (child.name or "") .. " is not expandable, skipping")
         end
       end
     end
@@ -53,16 +54,18 @@ end
 --- async method
 ---@param state table current state of the source
 ---@param node table a node to expand
----@param prefetcher table an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
+---@param prefetcher table? an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
 M.expand_directory_recursively = function(state, node, prefetcher)
   log.debug("Expanding directory " .. node:get_id())
-  if node.type ~= "directory" then
+  prefetcher = prefetcher or M.default_prefetcher
+  if not utils.is_expandable(node) then
     return
   end
-  state.explicitly_opened_directories = state.explicitly_opened_directories or {}
+
+  state.explicitly_opened_nodes = state.explicitly_opened_nodes or {}
   if prefetcher.should_prefetch(node) then
     local id = node:get_id()
-    state.explicitly_opened_directories[id] = true
+    state.explicitly_opened_nodes[id] = true
     prefetcher.prefetch(state, node)
     expand_loaded(node, state, prefetcher)
   else

Diff for: lua/neo-tree/sources/filesystem/commands.lua

@@ -68,9 +68,6 @@ M.delete_visual = function(state, selected_nodes)
 end
 
 M.expand_all_nodes = function(state, node)
-  if node == nil then
-    node = state.tree:get_node(state.path)
-  end
   cc.expand_all_nodes(state, node, fs.prefetcher)
 end
 

Diff for: lua/neo-tree/sources/filesystem/init.lua

@@ -64,13 +64,13 @@ local follow_internal = function(callback, force_show, async)
 
   log.debug("follow file: ", path_to_reveal)
   local show_only_explicitly_opened = function()
-    state.explicitly_opened_directories = state.explicitly_opened_directories or {}
+    state.explicitly_opened_nodes = state.explicitly_opened_nodes or {}
     local expanded_nodes = renderer.get_expanded_nodes(state.tree)
     local state_changed = false
     for _, id in ipairs(expanded_nodes) do
-      if not state.explicitly_opened_directories[id] then
+      if not state.explicitly_opened_nodes[id] then
         if path_to_reveal:sub(1, #id) == id then
-          state.explicitly_opened_directories[id] = state.follow_current_file.leave_dirs_open
+          state.explicitly_opened_nodes[id] = state.follow_current_file.leave_dirs_open
         else
           local node = state.tree:get_node(id)
           if node then
@@ -394,20 +394,20 @@ M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursiv
   if node.type ~= "directory" then
     return
   end
-  state.explicitly_opened_directories = state.explicitly_opened_directories or {}
+  state.explicitly_opened_nodes = state.explicitly_opened_nodes or {}
   if node.loaded == false then
     local id = node:get_id()
-    state.explicitly_opened_directories[id] = true
+    state.explicitly_opened_nodes[id] = true
     renderer.position.set(state, nil)
     fs_scan.get_items(state, id, path_to_reveal, callback, false, recursive)
   elseif node:has_children() then
     local updated = false
     if node:is_expanded() then
       updated = node:collapse()
-      state.explicitly_opened_directories[node:get_id()] = false
+      state.explicitly_opened_nodes[node:get_id()] = false
     else
       updated = node:expand()
-      state.explicitly_opened_directories[node:get_id()] = true
+      state.explicitly_opened_nodes[node:get_id()] = true
     end
     if updated and not skip_redraw then
       renderer.redraw(state)