feat(picker.git_diff): git_diff now also shows staged hunks and added stage/unstage/restore actions for hunks. Closes #2382

lua/snacks/picker/actions.lua

@@ -344,15 +344,21 @@ function M.git_stage(picker)
   local items = picker:selected({ fallback = true })
   local done = 0
   for _, item in ipairs(items) do
+    local opts = { cwd = item.cwd } ---@type snacks.picker.util.cmd.Opts
+
     local cmd = item.status:sub(2) == " " and { "git", "restore", "--staged", item.file } or { "git", "add", item.file }
-    Snacks.picker.util.cmd(cmd, function(data, code)
+    if item.diff then
+      opts.input = item.diff
+      cmd = { "git", "apply", "--cached", item.staged and "--reverse" or nil }
+    end
+    Snacks.picker.util.cmd(cmd, function()
       done = done + 1
       if done == #items then
         picker.list:set_selected()
         picker.list:set_target()
         picker:find()
       end
-    end, { cwd = item.cwd })
+    end, opts)
   end
 end
 
@@ -370,14 +376,22 @@ function M.git_restore(picker)
   local msg = #items == 1 and ("Discard changes to `%s`?"):format(files[1])
     or ("Discard changes to %d files?"):format(#items)
 
-  Snacks.picker.select({ "No", "Yes" }, { prompt = msg }, function(_, idx)
-    if not idx and idx == 2 then
-      return
-    end
+  Snacks.picker.util.confirm(msg, function()
     local done = 0
     for _, item in ipairs(items) do
       local cmd = { "git", "restore", item.file }
-      Snacks.picker.util.cmd(cmd, function(data, code)
+      local opts = { cwd = item.cwd }
+
+      if item.diff then
+        opts.input = item.diff
+        if item.staged then
+          cmd = { "git", "apply", "--reverse", "--cached" }
+        else
+          cmd = { "git", "apply", "--reverse" }
+        end
+      end
+
+      Snacks.picker.util.cmd(cmd, function()
         done = done + 1
         if done == #items then
           vim.schedule(function()
@@ -385,9 +399,10 @@ function M.git_restore(picker)
             picker.list:set_target()
             picker:find()
             vim.cmd.startinsert()
+            vim.cmd.checktime()
           end)
         end
-      end, { cwd = item.cwd })
+      end, opts)
     end
   end)
 end

lua/snacks/picker/config/sources.lua

@@ -340,8 +340,17 @@ M.git_status = {
 M.git_diff = {
   group = false,
   finder = "git_diff",
-  format = "file",
+  format = "git_status",
   preview = "diff",
+  matcher = { sort_empty = true },
+  win = {
+    input = {
+      keys = {
+        ["<Tab>"] = { "git_stage", mode = { "n", "i" } },
+        ["<c-r>"] = { "git_restore", mode = { "n", "i" } },
+      },
+    },
+  },
 }
 
 ---@class snacks.picker.grep.Config: snacks.picker.proc.Config

lua/snacks/picker/format.lua

@@ -573,6 +573,7 @@ function M.git_status(item, picker)
     ["?"] = "SnacksPickerGitStatusUntracked",
   }
   local hl = hls[s] or "SnacksPickerGitStatus"
+  hl = item.status:sub(1, 1) == "M" and "SnacksPickerGitStatusStaged" or hl
   ret[#ret + 1] = { a(item.status, 2, { align = "right" }), hl }
   ret[#ret + 1] = { " " }
   if item.rename then

lua/snacks/picker/source/git.lua

@@ -263,7 +263,10 @@ function M.diff(opts, ctx)
   if opts.staged then
     table.insert(args, "--cached")
   end
-  return require("snacks.picker.source.diff").diff(
+
+  local Diff = require("snacks.picker.source.diff")
+  local finders = {} ---@type snacks.picker.finder.result[]
+  finders[#finders + 1] = Diff.diff(
     ctx:opts({
       cmd = "git",
       args = args,
@@ -271,6 +274,39 @@ function M.diff(opts, ctx)
     }),
     ctx
   )
+  if opts.staged == nil and opts.base == nil then
+    finders[#finders + 1] = Diff.diff(
+      ctx:opts({
+        cmd = "git",
+        args = vim.list_extend(vim.deepcopy(args), { "--cached" }),
+        cwd = ctx:git_root(),
+      }),
+      ctx
+    )
+  end
+  return function(cb)
+    local items = {} ---@type snacks.picker.finder.Item[]
+    for f, finder in ipairs(finders) do
+      finder(function(item)
+        item.staged = opts.staged or f == 2
+        if item.staged then
+          item.status = "M "
+        else
+          item.status = " M"
+        end
+        items[#items + 1] = item
+      end)
+    end
+    table.sort(items, function(a, b)
+      if a.file ~= b.file then
+        return a.file < b.file
+      end
+      return a.pos[1] < b.pos[1]
+    end)
+    for _, item in ipairs(items) do
+      cb(item)
+    end
+  end
 end
 
 ---@param opts snacks.picker.git.branches.Config

lua/snacks/picker/util/init.lua

@@ -83,10 +83,12 @@ function M.confirm(prompt, fn)
   end)
 end
 
+---@alias snacks.picker.util.cmd.Opts {env?: table<string, string>, cwd?: string, input?: string}
 ---@param cmd string|string[]
 ---@param cb fun(output: string[], code: number)
----@param opts? {env?: table<string, string>, cwd?: string}
+---@param opts? snacks.picker.util.cmd.Opts
 function M.cmd(cmd, cb, opts)
+  opts = opts or {}
   local output = {} ---@type string[]
   local id = vim.fn.jobstart(
     cmd,
@@ -114,6 +116,9 @@ function M.cmd(cmd, cb, opts)
   )
   if id <= 0 then
     Snacks.notify.error(("Failed to start job `%s`"):format(cmd))
+  elseif opts.input then
+    vim.fn.chansend(id, opts.input .. "\n")
+    vim.fn.chanclose(id, "stdin")
   end
   return id > 0 and id or nil
 end