Merge pull request #21402 from lewis6991/feat/fs_ls

4 files changed, 133 insertions, 9 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index 9c98ed7771..1459392a81 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -2295,13 +2295,18 @@ basename({file})                                           *vim.fs.basename()*
     Return: ~
         (string) Basename of {file}
 
-dir({path})                                                     *vim.fs.dir()*
+dir({path}, {opts})                                             *vim.fs.dir()*
     Return an iterator over the files and directories located in {path}
 
     Parameters: ~
       • {path}  (string) An absolute or relative path to the directory to
                 iterate over. The path is first normalized
                 |vim.fs.normalize()|.
+      • {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
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 013089acc7..2ce0bd4de2 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -77,6 +77,9 @@ The following new APIs or features were added.
   Similarly, the `virtual_text` configuration in |vim.diagnostic.config()| now
   has a `suffix` option which does nothing by default.
 
+• |vim.fs.dir()| now has a `opts` argument with a depth field to allow
+  recursively searching a directory tree.
+
 • |vim.secure.read()| reads a file and prompts the user if it should be
   trusted and, if so, returns the file's contents.
 
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('\\', '/')
diff --git a/test/functional/lua/fs_spec.lua b/test/functional/lua/fs_spec.lua
index 88ad6ba24a..5db92ca174 100644
--- a/test/functional/lua/fs_spec.lua
+++ b/test/functional/lua/fs_spec.lua
@@ -141,6 +141,74 @@ describe('vim.fs', function()
         return false
       ]], nvim_dir, nvim_prog_basename))
     end)
+
+    it('works with opts.depth and opts.skip', function()
+      helpers.funcs.system 'mkdir -p testd/a/b/c'
+      helpers.funcs.system('touch '..table.concat({
+        'testd/a1',
+        'testd/b1',
+        'testd/c1',
+        'testd/a/a2',
+        'testd/a/b2',
+        'testd/a/c2',
+        'testd/a/b/a3',
+        'testd/a/b/b3',
+        'testd/a/b/c3',
+        'testd/a/b/c/a4',
+        'testd/a/b/c/b4',
+        'testd/a/b/c/c4',
+      }, ' '))
+
+      local function run(dir, depth, skip)
+         local r = exec_lua([[
+          local dir, depth, skip = ...
+          local r = {}
+          local skip_f
+          if skip then
+            skip_f = function(n)
+              if vim.tbl_contains(skip or {}, n) then
+                return false
+              end
+            end
+          end
+          for name, type_ in vim.fs.dir(dir, { depth = depth, skip = skip_f }) do
+            r[name] = type_
+          end
+          return r
+        ]], dir, depth, skip)
+        return r
+      end
+
+      local exp = {}
+
+      exp['a1'] = 'file'
+      exp['b1'] = 'file'
+      exp['c1'] = 'file'
+      exp['a'] = 'directory'
+
+      eq(exp, run('testd', 1))
+
+      exp['a/a2'] = 'file'
+      exp['a/b2'] = 'file'
+      exp['a/c2'] = 'file'
+      exp['a/b'] = 'directory'
+
+      eq(exp, run('testd', 2))
+
+      exp['a/b/a3'] = 'file'
+      exp['a/b/b3'] = 'file'
+      exp['a/b/c3'] = 'file'
+      exp['a/b/c'] = 'directory'
+
+      eq(exp, run('testd', 3))
+      eq(exp, run('testd', 999, {'a/b/c'}))
+
+      exp['a/b/c/a4'] = 'file'
+      exp['a/b/c/b4'] = 'file'
+      exp['a/b/c/c4'] = 'file'
+
+      eq(exp, run('testd', 999))
+    end)
   end)
 
   describe('find()', function()