1 files changed, 12 insertions, 3 deletions
diff --git a/lua/telescope/actions/init.lua b/lua/telescope/actions/init.lua
index 527e0b5..4b152b6 100644
--- a/lua/telescope/actions/init.lua
+++ b/lua/telescope/actions/init.lua
@@ -243,11 +243,20 @@ actions.git_checkout = function(prompt_bufnr)
   os.execute('git checkout ' .. val)
 end
 
-actions.git_add = function(prompt_bufnr)
+actions.git_staging_toggle = function(prompt_bufnr)
   local selection = actions.get_selected_entry(prompt_bufnr)
+
+  -- If parts of the file are staged and unstaged at the same time, stage
+  -- changes. Else toggle between staged and unstaged if the file is tracked,
+  -- and between added and untracked if the file is untracked.
+  if selection.status:sub(2) == ' ' then
+    os.execute('git restore --staged ' .. selection.value)
+  else
+    os.execute('git add ' .. selection.value)
+  end
   actions.close(prompt_bufnr)
-  local val = selection.value
-  os.execute('git add ' .. val)
+  require('telescope.builtin').git_status()
+  vim.api.nvim_feedkeys('i', 'n', false)
 end
 
 -- ==================================================