feat(fs): add opts argument to vim.fs.dir()

Added option depth to allow recursively searching a directory tree.

1 files changed, 56 insertions, 8 deletions

diff --git a/runtime/lua/vim/fs.lua b/runtime/lua/vim/fs.lua
index fdf3d29e94..89a1d0d345 100644
--- a/runtime/lua/vim/fs.lua
+++ b/runtime/lua/vim/fs.lua
@@ -72,18 +72,65 @@ function M.basename(file)
   return file:match('[/\\]$') and '' or (file:match('[^\\/]*$'):gsub('\\', '/'))
 end
 
+---@private
+local function join_paths(...)
+  return table.concat({ ... }, '/')
+end
+
 --- Return an iterator over the files and directories located in {path}
 ---
 ---@param path (string) An absolute or relative path to the directory to iterate
 ---            over. The path is first normalized |vim.fs.normalize()|.
+--- @param opts table|nil Optional keyword arguments:
+---             - depth: integer|nil How deep the traverse (default 1)
+---             - skip: (fun(dir_name: string): boolean)|nil Predicate
+---               to control traversal. Return false to stop searching the current directory.
+---               Only useful when depth > 1
+---
 ---@return Iterator over files and directories in {path}. Each iteration yields
 ---        two values: name and type. Each "name" is the basename of the file or
 ---        directory relative to {path}. Type is one of "file" or "directory".
-function M.dir(path)
-  return function(fs)
-    return vim.loop.fs_scandir_next(fs)
-  end,
-    vim.loop.fs_scandir(M.normalize(path))
+function M.dir(path, opts)
+  opts = opts or {}
+
+  vim.validate({
+    path = { path, { 'string' } },
+    depth = { opts.depth, { 'number' }, true },
+    skip = { opts.skip, { 'function' }, true },
+  })
+
+  if not opts.depth or opts.depth == 1 then
+    return function(fs)
+      return vim.loop.fs_scandir_next(fs)
+    end,
+      vim.loop.fs_scandir(M.normalize(path))
+  end
+
+  --- @async
+  return coroutine.wrap(function()
+    local dirs = { { path, 1 } }
+    while #dirs > 0 do
+      local dir0, level = unpack(table.remove(dirs, 1))
+      local dir = level == 1 and dir0 or join_paths(path, dir0)
+      local fs = vim.loop.fs_scandir(M.normalize(dir))
+      while fs do
+        local name, t = vim.loop.fs_scandir_next(fs)
+        if not name then
+          break
+        end
+        local f = level == 1 and name or join_paths(dir0, name)
+        coroutine.yield(f, t)
+        if
+          opts.depth
+          and level < opts.depth
+          and t == 'directory'
+          and (not opts.skip or opts.skip(f) ~= false)
+        then
+          dirs[#dirs + 1] = { f, level + 1 }
+        end
+      end
+    end
+  end)
 end
 
 --- Find files or directories in the given path.
@@ -155,7 +202,7 @@ function M.find(names, opts)
         local t = {}
         for name, type in M.dir(p) do
           if names(name) and (not opts.type or opts.type == type) then
-            table.insert(t, p .. '/' .. name)
+            table.insert(t, join_paths(p, name))
           end
         end
         return t
@@ -164,7 +211,7 @@ function M.find(names, opts)
       test = function(p)
         local t = {}
         for _, name in ipairs(names) do
-          local f = p .. '/' .. name
+          local f = join_paths(p, name)
           local stat = vim.loop.fs_stat(f)
           if stat and (not opts.type or opts.type == stat.type) then
             t[#t + 1] = f
@@ -201,7 +248,7 @@ function M.find(names, opts)
       end
 
       for other, type_ in M.dir(dir) do
-        local f = dir .. '/' .. other
+        local f = join_paths(dir, other)
         if type(names) == 'function' then
           if names(other) and (not opts.type or opts.type == type_) then
             if add(f) then
@@ -251,6 +298,7 @@ function M.normalize(path)
   vim.validate({ path = { path, 's' } })
   return (
     path
+      :gsub('^~$', vim.loop.os_homedir())
       :gsub('^~/', vim.loop.os_homedir() .. '/')
       :gsub('%$([%w_]+)', vim.loop.os_getenv)
       :gsub('\\', '/')