feat (file_selector) Add directory selection support to file selector (#954)

Co-authored-by: yetone <yetoneful@gmail.com>

lua/avante/file_selector.lua

@@ -11,18 +11,70 @@ local FileSelector = {}
 --- @class FileSelector
 --- @field id integer
 --- @field selected_filepaths string[]
---- @field file_cache string[]
 --- @field event_handlers table<string, function[]>
 
 ---@alias FileSelectorHandler fun(self: FileSelector, on_select: fun(filepaths: string[] | nil)): nil
 
+function FileSelector:process_directory(absolute_path, project_root)
+  local files = scan.scan_dir(absolute_path, {
+    hidden = false,
+    depth = math.huge,
+    add_dirs = false,
+    respect_gitignore = true,
+  })
+
+  for _, file in ipairs(files) do
+    local rel_path = Path:new(file):make_relative(project_root)
+    if not vim.tbl_contains(self.selected_filepaths, rel_path) then table.insert(self.selected_filepaths, rel_path) end
+  end
+  self:emit("update")
+end
+
+---@param selected_paths string[] | nil
+---@return nil
+function FileSelector:handle_path_selection(selected_paths)
+  if not selected_paths then return end
+  local project_root = Utils.get_project_root()
+
+  for _, selected_path in ipairs(selected_paths) do
+    local absolute_path = Path:new(project_root):joinpath(selected_path):absolute()
+
+    local stat = vim.loop.fs_stat(absolute_path)
+    if stat and stat.type == "directory" then
+      self.process_directory(self, absolute_path, project_root)
+    else
+      local uniform_path = Utils.uniform_path(selected_path)
+      if Config.file_selector.provider == "native" then
+        table.insert(self.selected_filepaths, uniform_path)
+      else
+        if not vim.tbl_contains(self.selected_filepaths, uniform_path) then
+          table.insert(self.selected_filepaths, uniform_path)
+        end
+      end
+    end
+  end
+  self:emit("update")
+end
+
+local function get_project_filepaths()
+  local project_root = Utils.get_project_root()
+  local files = Utils.scan_directory_respect_gitignore({ directory = project_root, add_dirs = true })
+  files = vim.iter(files):map(function(filepath) return Path:new(filepath):make_relative(project_root) end):totable()
+
+  return vim.tbl_map(function(path)
+    local rel_path = Path:new(path):make_relative(project_root)
+    local stat = vim.loop.fs_stat(path)
+    if stat and stat.type == "directory" then rel_path = rel_path .. "/" end
+    return rel_path
+  end, files)
+end
+
 ---@param id integer
 ---@return FileSelector
 function FileSelector:new(id)
   return setmetatable({
     id = id,
     selected_files = {},
-    file_cache = {},
     event_handlers = {},
   }, { __index = self })
 end
@@ -35,6 +87,13 @@ end
 function FileSelector:add_selected_file(filepath)
   if not filepath or filepath == "" then return end
 
+  local absolute_path = Path:new(Utils.get_project_root()):joinpath(filepath):absolute()
+  local stat = vim.loop.fs_stat(absolute_path)
+
+  if stat and stat.type == "directory" then
+    self.process_directory(self, absolute_path, Utils.get_project_root())
+    return
+  end
   local uniform_path = Utils.uniform_path(filepath)
 
   -- Avoid duplicates
@@ -104,26 +163,29 @@ function FileSelector:off(event, callback)
   end
 end
 
----@return nil
+function FileSelector:open() self:show_select_ui() end
-function FileSelector:open()
-  if Config.file_selector.provider == "native" then self:update_file_cache() end
-  self:show_select_ui()
-end
 
----@return nil
+function FileSelector:get_filepaths()
-function FileSelector:update_file_cache()
+  local filepaths = get_project_filepaths()
-  local project_root = Path:new(Utils.get_project_root()):absolute()
 
-  local filepaths = scan.scan_dir(project_root, {
+  table.sort(filepaths, function(a, b)
-    respect_gitignore = true,
+    local a_stat = vim.loop.fs_stat(a)
-  })
+    local b_stat = vim.loop.fs_stat(b)
+    local a_is_dir = a_stat and a_stat.type == "directory"
+    local b_is_dir = b_stat and b_stat.type == "directory"
 
-  -- Sort buffer names alphabetically
+    if a_is_dir and not b_is_dir then
-  table.sort(filepaths, function(a, b) return a < b end)
+      return true
+    elseif not a_is_dir and b_is_dir then
+      return false
+    else
+      return a < b
+    end
+  end)
 
-  self.file_cache = vim
+  return vim
     .iter(filepaths)
-    :map(function(filepath) return Path:new(filepath):make_relative(project_root) end)
+    :filter(function(filepath) return not vim.tbl_contains(self.selected_filepaths, filepath) end)
     :totable()
 end
 
@@ -135,28 +197,32 @@ function FileSelector:fzf_ui(handler)
     return
   end
 
-  local close_action = function() handler(nil) end
+  local filepaths = self:get_filepaths()
-  fzf_lua.files(vim.tbl_deep_extend("force", {
-    file_ignore_patterns = self.selected_filepaths,
-    prompt = string.format("%s> ", PROMPT_TITLE),
-    fzf_opts = {},
-    git_icons = false,
-    actions = {
-      ["default"] = function(selected)
-        if not selected or #selected == 0 then return close_action() end
-        ---@type string[]
-        local selections = {}
-        for _, entry in ipairs(selected) do
-          local file = fzf_lua.path.entry_to_file(entry)
-          if file and file.path then table.insert(selections, file.path) end
-        end
 
-        handler(selections)
+  local close_action = function() handler(nil) end
-      end,
+  fzf_lua.fzf_exec(
-      ["esc"] = close_action,
+    filepaths,
-      ["ctrl-c"] = close_action,
+    vim.tbl_deep_extend("force", {
-    },
+      prompt = string.format("%s> ", PROMPT_TITLE),
-  }, Config.file_selector.provider_opts))
+      fzf_opts = {},
+      git_icons = false,
+      actions = {
+        ["default"] = function(selected)
+          if not selected or #selected == 0 then return close_action() end
+          ---@type string[]
+          local selections = {}
+          for _, entry in ipairs(selected) do
+            local file = fzf_lua.path.entry_to_file(entry)
+            if file and file.path then table.insert(selections, file.path) end
+          end
+
+          handler(selections)
+        end,
+        ["esc"] = close_action,
+        ["ctrl-c"] = close_action,
+      },
+    }, Config.file_selector.provider_opts)
+  )
 end
 
 function FileSelector:mini_pick_ui(handler)
@@ -194,9 +260,7 @@ function FileSelector:telescope_ui(handler)
   local action_state = require("telescope.actions.state")
   local action_utils = require("telescope.actions.utils")
 
-  local project_root = Utils.get_project_root()
+  local files = self:get_filepaths()
-  local files = Utils.scan_directory_respect_gitignore(project_root)
-  files = vim.iter(files):map(function(filepath) return Path:new(filepath):make_relative(project_root) end):totable()
 
   pickers
     .new(
@@ -208,7 +272,6 @@ function FileSelector:telescope_ui(handler)
         sorter = conf.file_sorter(),
         attach_mappings = function(prompt_bufnr, map)
           map("i", "<esc>", require("telescope.actions").close)
-
           actions.select_default:replace(function()
             local picker = action_state.get_current_picker(prompt_bufnr)
 
@@ -234,10 +297,7 @@ function FileSelector:telescope_ui(handler)
 end
 
 function FileSelector:native_ui(handler)
-  local filepaths = vim
+  local filepaths = self:get_filepaths()
-    .iter(self.file_cache)
-    :filter(function(filepath) return not vim.tbl_contains(self.selected_filepaths, filepath) end)
-    :totable()
 
   vim.ui.select(filepaths, {
     prompt = string.format("%s:", PROMPT_TITLE),
@@ -253,24 +313,7 @@ end
 
 ---@return nil
 function FileSelector:show_select_ui()
-  ---@param filepaths string[] | nil
+  local function handler(selected_paths) self:handle_path_selection(selected_paths) end
-  ---@return nil
-  local handler = function(filepaths)
-    if not filepaths then return end
-
-    for _, filepath in ipairs(filepaths) do
-      local uniform_path = Utils.uniform_path(filepath)
-      if Config.file_selector.provider == "native" then
-        table.insert(self.selected_filepaths, uniform_path)
-      else
-        if not vim.tbl_contains(self.selected_filepaths, uniform_path) then
-          table.insert(self.selected_filepaths, uniform_path)
-        end
-      end
-    end
-
-    self:emit("update")
-  end
 
   vim.schedule(function()
     if Config.file_selector.provider == "native" then

lua/avante/repo_map.lua

@@ -52,7 +52,11 @@ function RepoMap._build_repo_map(project_root, file_ext)
   local ignore_patterns = vim.list_extend(gitignore_patterns, Config.repo_map.ignore_patterns)
   local negate_patterns = vim.list_extend(gitignore_negate_patterns, Config.repo_map.negate_patterns)
 
-  local filepaths = Utils.scan_directory(project_root, ignore_patterns, negate_patterns)
+  local filepaths = Utils.scan_directory({
+    directory = project_root,
+    gitignore_patterns = ignore_patterns,
+    gitignore_negate_patterns = negate_patterns,
+  })
   if filepaths and not RepoMap._init_repo_map_lib() then
     -- or just throw an error if we don't want to execute request without codebase
     Utils.error("Failed to load avante_repo_map")

lua/avante/utils/init.lua

@@ -259,15 +259,7 @@ end
 
 ---@param path string
 ---@return string
-function M.norm(path)
+function M.norm(path) return vim.fs.normalize(path) end
-  if path:sub(1, 1) == "~" then
-    local home = vim.uv.os_homedir()
-    if home:sub(-1) == "\\" or home:sub(-1) == "/" then home = home:sub(1, -2) end
-    path = home .. path:sub(2)
-  end
-  path = path:gsub("\\", "/"):gsub("/+", "/")
-  return path:sub(-1) == "/" and path:sub(1, -2) or path
-end
 
 ---@param msg string|string[]
 ---@param opts? LazyNotifyOpts
@@ -643,14 +635,27 @@ function M.is_ignored(file, ignore_patterns, negate_patterns)
   return false
 end
 
-function M.scan_directory_respect_gitignore(directory)
+---@param options { directory: string, add_dirs?: boolean }
+function M.scan_directory_respect_gitignore(options)
+  local directory = options.directory
   local gitignore_path = directory .. "/.gitignore"
   local gitignore_patterns, gitignore_negate_patterns = M.parse_gitignore(gitignore_path)
   gitignore_patterns = vim.list_extend(gitignore_patterns, { "%.git", "%.worktree", "__pycache__", "node_modules" })
-  return M.scan_directory(directory, gitignore_patterns, gitignore_negate_patterns)
+  return M.scan_directory({
+    directory = directory,
+    gitignore_patterns = gitignore_patterns,
+    gitignore_negate_patterns = gitignore_negate_patterns,
+    add_dirs = options.add_dirs,
+  })
 end
 
-function M.scan_directory(directory, ignore_patterns, negate_patterns)
+---@param options { directory: string, gitignore_patterns: string[], gitignore_negate_patterns: string[], add_dirs?: boolean }
+function M.scan_directory(options)
+  local directory = options.directory
+  local ignore_patterns = options.gitignore_patterns
+  local negate_patterns = options.gitignore_negate_patterns
+  local add_dirs = options.add_dirs or false
+
   local files = {}
   local handle = vim.loop.fs_scandir(directory)
 
@@ -662,7 +667,18 @@ function M.scan_directory(directory, ignore_patterns, negate_patterns)
 
     local full_path = directory .. "/" .. name
     if type == "directory" then
-      vim.list_extend(files, M.scan_directory(full_path, ignore_patterns, negate_patterns))
+      if add_dirs and not M.is_ignored(full_path, ignore_patterns, negate_patterns) then
+        table.insert(files, full_path)
+      end
+      vim.list_extend(
+        files,
+        M.scan_directory({
+          directory = full_path,
+          gitignore_patterns = ignore_patterns,
+          gitignore_negate_patterns = negate_patterns,
+          add_dirs = add_dirs,
+        })
+      )
     elseif type == "file" then
       if not M.is_ignored(full_path, ignore_patterns, negate_patterns) then table.insert(files, full_path) end
     end