feat(#3132): add api.node.expand and api.node.collapse (#3133)

* feat: allow passing node to collapse all

* refactor: use snake case

* feat: handle api legacy calls and update signature

* refactor: make sure open is a boolean

* doc: collapse_all

* Revert "doc: collapse_all"

This reverts commit d243da3e14.

* add api.node.collapse

* add api.node.expand

* add api.node.expand

---------

Co-authored-by: Alexander Courtis <alex@courtis.org>

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

@@ -2,6 +2,7 @@ local utils = require("nvim-tree.utils")
 local core = require("nvim-tree.core")
 local Iterator = require("nvim-tree.iterators.node-iterator")
 
+local FileNode = require("nvim-tree.node.file")
 local DirectoryNode = require("nvim-tree.node.directory")
 
 local M = {}
@@ -23,26 +24,30 @@ local function buf_match()
   end
 end
 
----@param keep_buffers boolean
-function M.fn(keep_buffers)
+---Collapse a node, root if nil
+---@param node Node?
+---@param opts ApiCollapseOpts
+local function collapse(node, opts)
   local explorer = core.get_explorer()
   if not explorer then
     return
   end
 
-  local node = explorer:get_node_at_cursor()
-  if not node then
+  node = node or explorer
+
+  local node_at_cursor = explorer:get_node_at_cursor()
+  if not node_at_cursor then
     return
   end
 
   local matches = buf_match()
 
-  Iterator.builder(explorer.nodes)
+  Iterator.builder({ node:is(FileNode) and node.parent or node:as(DirectoryNode) })
     :hidden()
     :applier(function(n)
       local dir = n:as(DirectoryNode)
       if dir then
-        dir.open = keep_buffers and matches(dir.absolute_path)
+        dir.open = opts.keep_buffers == true and matches(dir.absolute_path)
       end
     end)
     :recursor(function(n)
@@ -51,7 +56,26 @@ function M.fn(keep_buffers)
     :iterate()
 
   explorer.renderer:draw()
-  utils.focus_node_or_parent(node)
+  utils.focus_node_or_parent(node_at_cursor)
+end
+
+
+---@param opts ApiCollapseOpts|boolean|nil legacy -> opts.keep_buffers
+function M.all(opts)
+  -- legacy arguments
+  if type(opts) == "boolean" then
+    opts = {
+      keep_buffers = opts,
+    }
+  end
+
+  collapse(nil, opts or {})
+end
+
+---@param node Node
+---@param opts ApiCollapseOpts?
+function M.node(node, opts)
+  collapse(node, opts or {})
 end
 
 return M

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

@@ -2,6 +2,7 @@ local core = require("nvim-tree.core")
 local Iterator = require("nvim-tree.iterators.node-iterator")
 local notify = require("nvim-tree.notify")
 
+local FileNode = require("nvim-tree.node.file")
 local DirectoryNode = require("nvim-tree.node.directory")
 
 local M = {}
@@ -70,23 +71,38 @@ 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()
-  local parent = node:as(DirectoryNode) or explorer
-  if not parent then
+---@param node Node?
+local function expand_node(node)
+  if not node then
     return
   end
 
-  if gen_iterator()(parent) then
+  if gen_iterator()(node) then
     notify.warn("expansion iteration was halted after " .. M.MAX_FOLDER_DISCOVERY .. " discovered folders")
   end
+
+  local explorer = core.get_explorer()
   if explorer then
     explorer.renderer:draw()
   end
 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())
+end
+
+---Expand the directory node or parent node
+---@param node Node
+function M.node(node)
+  if not node then
+    return
+  end
+
+  expand_node(node:is(FileNode) and node.parent or node:as(DirectoryNode))
+end
+
 function M.setup(opts)
   M.MAX_FOLDER_DISCOVERY = opts.actions.expand_all.max_folder_discovery
   M.EXCLUDE = to_lookup_table(opts.actions.expand_all.exclude)

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

@@ -1,10 +1,10 @@
 local M = {}
 
-M.collapse_all = require("nvim-tree.actions.tree.modifiers.collapse-all")
-M.expand_all = require("nvim-tree.actions.tree.modifiers.expand-all")
+M.collapse = require("nvim-tree.actions.tree.modifiers.collapse")
+M.expand = require("nvim-tree.actions.tree.modifiers.expand")
 
 function M.setup(opts)
-  M.expand_all.setup(opts)
+  M.expand.setup(opts)
 end
 
 return M