From 77fb3fd81e78afd96ef017b2ad9107c1775a78b5 Mon Sep 17 00:00:00 2001
From: Andy Stewart <boss@airbladesoftware.com>
Date: Wed, 21 Aug 2019 10:27:23 +0100
Subject: [PATCH] Enable :w in preview window to stage the hunk.

---
 README.mkd                  |  2 +-
 autoload/gitgutter/hunk.vim |  4 ++++
 doc/gitgutter.txt           |  3 ++-
 test/test_gitgutter.vim     | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/README.mkd b/README.mkd
index f4616e9..3261f90 100644
--- a/README.mkd
+++ b/README.mkd
@@ -198,7 +198,7 @@ To stage part of any hunk:
 * preview the hunk, e.g. `<Leader>hp`;
 * move to the preview window, e.g. `:wincmd P`;
 * delete the lines you do not want to stage;
-* stage the remaining lines, e.g. `<Leader>hs` or `:GitGutterStageHunk`.
+* stage the remaining lines: either write (`:w`) the window or stage via `<Leader>hs` or `:GitGutterStageHunk`.
 
 See the FAQ if you want to unstage staged changes.
 
diff --git a/autoload/gitgutter/hunk.vim b/autoload/gitgutter/hunk.vim
index b35178f..6f4a6ff 100644
--- a/autoload/gitgutter/hunk.vim
+++ b/autoload/gitgutter/hunk.vim
@@ -315,6 +315,10 @@ function! s:preview(hunk_diff)
   call setline(1, body)
   normal! gg
 
+  cnoreabbrev <buffer> <expr> w  getcmdtype() == ':' && getcmdline() == 'w'  ? 'GitGutterStageHunk' : 'w'
+  " Staging hunk from the preview window closes the window anyway.
+  cnoreabbrev <buffer> <expr> wq getcmdtype() == ':' && getcmdline() == 'wq' ? 'GitGutterStageHunk' : 'wq'
+
   noautocmd wincmd p
 endfunction
 
diff --git a/doc/gitgutter.txt b/doc/gitgutter.txt
index 9e126b6..5dafcbb 100644
--- a/doc/gitgutter.txt
+++ b/doc/gitgutter.txt
@@ -170,7 +170,8 @@ Commands for operating on a hunk:~
 
                         To stage part of any hunk, first |GitGutterPreviewHunk|
                         it, then move to the preview window, delete the lines
-                        you do not want to stage, and |GitGutterStageHunk|.
+                        you do not want to stage, and |write| or
+                        |GitGutterStageHunk|.
 
                                                  *gitgutter-:GitGutterUndoHunk*
 :GitGutterUndoHunk      Undo the hunk the cursor is in.
diff --git a/test/test_gitgutter.vim b/test/test_gitgutter.vim
index 58c0565..9b0f48f 100644
--- a/test/test_gitgutter.vim
+++ b/test/test_gitgutter.vim
@@ -548,6 +548,42 @@ function Test_hunk_stage_partial_preview_added()
 endfunction
 
 
+function Test_hunk_stage_preview_write()
+  call append(5, ['A','B','C','D'])
+  6
+  GitGutterPreviewHunk
+  wincmd P
+
+  " preview window
+  w
+  " original window
+  write
+
+  let expected = [
+        \ {'lnum': 6, 'name': 'GitGutterLineAdded'},
+        \ {'lnum': 7, 'name': 'GitGutterLineAdded'},
+        \ {'lnum': 8, 'name': 'GitGutterLineAdded'},
+        \ {'lnum': 9, 'name': 'GitGutterLineAdded'},
+        \ ]
+  call s:assert_signs(expected, 'fixture.txt')
+
+  call assert_equal([], s:git_diff())
+
+  let expected = [
+        \ 'diff --git a/fixture.txt b/fixture.txt',
+        \ 'index f5c6aff..975852f 100644',
+        \ '--- a/fixture.txt',
+        \ '+++ b/fixture.txt',
+        \ '@@ -5,0 +6,4 @@ e',
+        \ '+A',
+        \ '+B',
+        \ '+C',
+        \ '+D',
+        \ ]
+  call assert_equal(expected, s:git_diff_staged())
+endfunction
+
+
 function Test_hunk_stage_partial_preview_added_removed()
   4,5delete
   call append(3, ['A','B','C','D'])