You can now stage changes one line at a time

Staging Lines

26 files changed, 1204 insertions, 8 deletions

diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index adc26b8c8..1d8e5a10d 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -485,7 +485,6 @@ func (c *GitCommand) getMergeBase() (string, error) {
 	output, err := c.OSCommand.RunCommandWithOutput(fmt.Sprintf("git merge-base HEAD %s", baseBranch))
 	if err != nil {
 		// swallowing error because it's not a big deal; probably because there are no commits yet
-		c.Log.Error(err)
 	}
 	return output, nil
 }
@@ -571,9 +570,10 @@ func (c *GitCommand) CheckRemoteBranchExists(branch *Branch) bool {
 }
 
 // Diff returns the diff of a file
-func (c *GitCommand) Diff(file *File) string {
+func (c *GitCommand) Diff(file *File, plain bool) string {
 	cachedArg := ""
 	trackedArg := "--"
+	colorArg := "--color"
 	fileName := c.OSCommand.Quote(file.Name)
 	if file.HasStagedChanges && !file.HasUnstagedChanges {
 		cachedArg = "--cached"
@@ -581,9 +581,25 @@ func (c *GitCommand) Diff(file *File) string {
 	if !file.Tracked && !file.HasStagedChanges {
 		trackedArg = "--no-index /dev/null"
 	}
-	command := fmt.Sprintf("git diff --color %s %s %s", cachedArg, trackedArg, fileName)
+	if plain {
+		colorArg = ""
+	}
+
+	command := fmt.Sprintf("git diff %s %s %s %s", colorArg, cachedArg, trackedArg, fileName)
 
 	// for now we assume an error means the file was deleted
 	s, _ := c.OSCommand.RunCommandWithOutput(command)
 	return s
 }
+
+func (c *GitCommand) ApplyPatch(patch string) (string, error) {
+	filename, err := c.OSCommand.CreateTempFile("patch", patch)
+	if err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+
+	defer func() { _ = c.OSCommand.RemoveFile(filename) }()
+
+	return c.OSCommand.RunCommandWithOutput(fmt.Sprintf("git apply --cached %s", filename))
+}
diff --git a/pkg/commands/git_test.go b/pkg/commands/git_test.go
index aec928ea5..206696b22 100644
--- a/pkg/commands/git_test.go
+++ b/pkg/commands/git_test.go
@@ -1770,6 +1770,7 @@ func TestGitCommandDiff(t *testing.T) {
 		testName string
 		command  func(string, ...string) *exec.Cmd
 		file     *File
+		plain    bool
 	}
 
 	scenarios := []scenario{
@@ -1786,6 +1787,22 @@ func TestGitCommandDiff(t *testing.T) {
 				HasStagedChanges: false,
 				Tracked:          true,
 			},
+			false,
+		},
+		{
+			"Default case",
+			func(cmd string, args ...string) *exec.Cmd {
+				assert.EqualValues(t, "git", cmd)
+				assert.EqualValues(t, []string{"diff", "--", "test.txt"}, args)
+
+				return exec.Command("echo")
+			},
+			&File{
+				Name:             "test.txt",
+				HasStagedChanges: false,
+				Tracked:          true,
+			},
+			true,
 		},
 		{
 			"All changes staged",
@@ -1801,6 +1818,7 @@ func TestGitCommandDiff(t *testing.T) {
 				HasUnstagedChanges: false,
 				Tracked:            true,
 			},
+			false,
 		},
 		{
 			"File not tracked and file has no staged changes",
@@ -1815,6 +1833,7 @@ func TestGitCommandDiff(t *testing.T) {
 				HasStagedChanges: false,
 				Tracked:          false,
 			},
+			false,
 		},
 	}
 
@@ -1822,7 +1841,7 @@ func TestGitCommandDiff(t *testing.T) {
 		t.Run(s.testName, func(t *testing.T) {
 			gitCmd := newDummyGitCommand()
 			gitCmd.OSCommand.command = s.command
-			gitCmd.Diff(s.file)
+			gitCmd.Diff(s.file, s.plain)
 		})
 	}
 }
@@ -1979,3 +1998,61 @@ func TestGitCommandCurrentBranchName(t *testing.T) {
 		})
 	}
 }
+
+func TestGitCommandApplyPatch(t *testing.T) {
+	type scenario struct {
+		testName string
+		command  func(string, ...string) *exec.Cmd
+		test     func(string, error)
+	}
+
+	scenarios := []scenario{
+		{
+			"valid case",
+			func(cmd string, args ...string) *exec.Cmd {
+				assert.Equal(t, "git", cmd)
+				assert.EqualValues(t, []string{"apply", "--cached"}, args[0:2])
+				filename := args[2]
+				content, err := ioutil.ReadFile(filename)
+				assert.NoError(t, err)
+
+				assert.Equal(t, "test", string(content))
+
+				return exec.Command("echo", "done")
+			},
+			func(output string, err error) {
+				assert.NoError(t, err)
+				assert.EqualValues(t, "done\n", output)
+			},
+		},
+		{
+			"command returns error",
+			func(cmd string, args ...string) *exec.Cmd {
+				assert.Equal(t, "git", cmd)
+				assert.EqualValues(t, []string{"apply", "--cached"}, args[0:2])
+				filename := args[2]
+				// TODO: Ideally we want to mock out OSCommand here so that we're not
+				// double handling testing it's CreateTempFile functionality,
+				// but it is going to take a bit of work to make a proper mock for it
+				// so I'm leaving it for another PR
+				content, err := ioutil.ReadFile(filename)
+				assert.NoError(t, err)
+
+				assert.Equal(t, "test", string(content))
+
+				return exec.Command("test")
+			},
+			func(output string, err error) {
+				assert.Error(t, err)
+			},
+		},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			gitCmd := newDummyGitCommand()
+			gitCmd.OSCommand.command = s.command
+			s.test(gitCmd.ApplyPatch("test"))
+		})
+	}
+}
diff --git a/pkg/commands/os.go b/pkg/commands/os.go
index 8b4b7879e..6b28a69bb 100644
--- a/pkg/commands/os.go
+++ b/pkg/commands/os.go
@@ -2,6 +2,7 @@ package commands
 
 import (
 	"errors"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"strings"
@@ -176,3 +177,28 @@ func (c *OSCommand) AppendLineToFile(filename, line string) error {
 	_, err = f.WriteString("\n" + line)
 	return err
 }
+
+// CreateTempFile writes a string to a new temp file and returns the file's name
+func (c *OSCommand) CreateTempFile(filename, content string) (string, error) {
+	tmpfile, err := ioutil.TempFile("", filename)
+	if err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+
+	if _, err := tmpfile.Write([]byte(content)); err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+	if err := tmpfile.Close(); err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+
+	return tmpfile.Name(), nil
+}
+
+// RemoveFile removes a file at the specified path
+func (c *OSCommand) RemoveFile(filename string) error {
+	return os.Remove(filename)
+}
diff --git a/pkg/commands/os_test.go b/pkg/commands/os_test.go
index ebb855cbe..a08c4b57d 100644
--- a/pkg/commands/os_test.go
+++ b/pkg/commands/os_test.go
@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"testing"
@@ -364,3 +365,34 @@ func TestOSCommandFileType(t *testing.T) {
 		_ = os.RemoveAll(s.path)
 	}
 }
+
+func TestOSCommandCreateTempFile(t *testing.T) {
+	type scenario struct {
+		testName string
+		filename string
+		content  string
+		test     func(string, error)
+	}
+
+	scenarios := []scenario{
+		{
+			"valid case",
+			"filename",
+			"content",
+			func(path string, err error) {
+				assert.NoError(t, err)
+
+				content, err := ioutil.ReadFile(path)
+				assert.NoError(t, err)
+
+				assert.Equal(t, "content", string(content))
+			},
+		},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			s.test(newDummyOSCommand().CreateTempFile(s.filename, s.content))
+		})
+	}
+}
diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
new file mode 100644
index 000000000..3c523232e
--- /dev/null
+++ b/pkg/git/patch_modifier.go
@@ -0,0 +1,156 @@
+package git
+
+import (
+	"errors"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/jesseduffield/lazygit/pkg/i18n"
+	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/sirupsen/logrus"
+)
+
+type PatchModifier struct {
+	Log *logrus.Entry
+	Tr  *i18n.Localizer
+}
+
+// NewPatchModifier builds a new branch list builder
+func NewPatchModifier(log *logrus.Entry) (*PatchModifier, error) {
+	return &PatchModifier{
+		Log: log,
+	}, nil
+}
+
+// ModifyPatchForHunk takes the original patch, which may contain several hunks,
+// and removes any hunks that aren't the selected hunk
+func (p *PatchModifier) ModifyPatchForHunk(patch string, hunkStarts []int, currentLine int) (string, error) {
+	// get hunk start and end
+	lines := strings.Split(patch, "\n")
+	hunkStartIndex := utils.PrevIndex(hunkStarts, currentLine)
+	hunkStart := hunkStarts[hunkStartIndex]
+	nextHunkStartIndex := utils.NextIndex(hunkStarts, currentLine)
+	var hunkEnd int
+	if nextHunkStartIndex == 0 {
+		hunkEnd = len(lines) - 1
+	} else {
+		hunkEnd = hunkStarts[nextHunkStartIndex]
+	}
+
+	headerLength, err := p.getHeaderLength(lines)
+	if err != nil {
+		return "", err
+	}
+
+	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+	output += strings.Join(lines[hunkStart:hunkEnd], "\n") + "\n"
+
+	return output, nil
+}
+
+func (p *PatchModifier) getHeaderLength(patchLines []string) (int, error) {
+	for index, line := range patchLines {
+		if strings.HasPrefix(line, "@@") {
+			return index, nil
+		}
+	}
+	return 0, errors.New(p.Tr.SLocalize("CantFindHunks"))
+}
+
+// ModifyPatchForLine takes the original patch, which may contain several hunks,
+// and the line number of the line we want to stage
+func (p *PatchModifier) ModifyPatchForLine(patch string, lineNumber int) (string, error) {
+	lines := strings.Split(patch, "\n")
+	headerLength, err := p.getHeaderLength(lines)
+	if err != nil {
+		return "", err
+	}
+	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+
+	hunkStart, err := p.getHunkStart(lines, lineNumber)
+	if err != nil {
+		return "", err
+	}
+
+	hunk, err := p.getModifiedHunk(lines, hunkStart, lineNumber)
+	if err != nil {
+		return "", err
+	}
+
+	output += strings.Join(hunk, "\n")
+
+	return output, nil
+}
+
+// getHunkStart returns the line number of the hunk we're going to be modifying
+// in order to stage our line
+func (p *PatchModifier) getHunkStart(patchLines []string, lineNumber int) (int, error) {
+	// find the hunk that we're modifying
+	hunkStart := 0
+	for index, line := range patchLines {
+		if strings.HasPrefix(line, "@@") {
+			hunkStart = index
+		}
+		if index == lineNumber {
+			return hunkStart, nil
+		}
+	}
+
+	return 0, errors.New(p.Tr.SLocalize("CantFindHunk"))
+}
+
+func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, lineNumber int) ([]string, error) {
+	lineChanges := 0
+	// strip the hunk down to just the line we want to stage
+	newHunk := []string{patchLines[hunkStart]}
+	for offsetIndex, line := range patchLines[hunkStart+1:] {
+		index := offsetIndex + hunkStart + 1
+		if strings.HasPrefix(line, "@@") {
+			newHunk = append(newHunk, "\n")
+			break
+		}
+		if index != lineNumber {
+			// we include other removals but treat them like context
+			if strings.HasPrefix(line, "-") {
+				newHunk = append(newHunk, " "+line[1:])
+				lineChanges += 1
+				continue
+			}
+			// we don't include other additions
+			if strings.HasPrefix(line, "+") {
+				lineChanges -= 1
+				continue
+			}
+		}
+		newHunk = append(newHunk, line)
+	}
+
+	var err error
+	newHunk[0], err = p.updatedHeader(newHunk[0], lineChanges)
+	if err != nil {
+		return nil, err
+	}
+
+	return newHunk, nil
+}
+
+// updatedHeader returns the hunk header with the updated line range
+// we need to update the hunk length to reflect the changes we made
+// if the hunk has three additions but we're only staging one, then
+// @@ -14,8 +14,11 @@ import (
+// becomes
+// @@ -14,8 +14,9 @@ import (
+func (p *PatchModifier) updatedHeader(currentHeader string, lineChanges int) (string, error) {
+	// current counter is the number after the second comma
+	re := regexp.MustCompile(`(\d+) @@`)
+	prevLengthString := re.FindStringSubmatch(currentHeader)[1]
+
+	prevLength, err := strconv.Atoi(prevLengthString)
+	if err != nil {
+		return "", err
+	}
+	re = regexp.MustCompile(`\d+ @@`)
+	newLength := strconv.Itoa(prevLength + lineChanges)
+	return re.ReplaceAllString(currentHeader, newLength+" @@"), nil
+}
diff --git a/pkg/git/patch_modifier_test.go b/pkg/git/patch_modifier_test.go
new file mode 100644
index 000000000..bc2073d55
--- /dev/null
+++ b/pkg/git/patch_modifier_test.go
@@ -0,0 +1,89 @@
+package git
+
+import (
+	"io/ioutil"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+)
+
+func newDummyLog() *logrus.Entry {
+	log := logrus.New()
+	log.Out = ioutil.Discard
+	return log.WithField("test", "test")
+}
+
+func newDummyPatchModifier() *PatchModifier {
+	return &PatchModifier{
+		Log: newDummyLog(),
+	}
+}
+func TestModifyPatchForLine(t *testing.T) {
+	type scenario struct {
+		testName              string
+		patchFilename         string
+		lineNumber            int
+		shouldError           bool
+		expectedPatchFilename string
+	}
+
+	scenarios := []scenario{
+		{
+			"Removing one line",
+			"testdata/testPatchBefore.diff",
+			8,
+			false,
+			"testdata/testPatchAfter1.diff",
+		},
+		{
+			"Adding one line",
+			"testdata/testPatchBefore.diff",
+			10,
+			false,
+			"testdata/testPatchAfter2.diff",
+		},
+		{
+			"Adding one line in top hunk in diff with multiple hunks",
+			"testdata/testPatchBefore2.diff",
+			20,
+			false,
+			"testdata/testPatchAfter3.diff",
+		},
+		{
+			"Adding one line in top hunk in diff with multiple hunks",
+			"testdata/testPatchBefore2.diff",
+			53,
+			false,
+			"testdata/testPatchAfter4.diff",
+		},
+		{
+			"adding unstaged file with a single line",
+			"testdata/addedFile.diff",
+			6,
+			false,
+			"testdata/addedFile.diff",
+		},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			p := newDummyPatchModifier()
+			beforePatch, err := ioutil.ReadFile(s.patchFilename)
+			if err != nil {
+				panic("Cannot open file at " + s.patchFilename)
+			}
+			afterPatch, err := p.ModifyPatchForLine(string(beforePatch), s.lineNumber)
+			if s.shouldError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				expected, err := ioutil.ReadFile(s.expectedPatchFilename)
+				if err != nil {
+					panic("Cannot open file at " + s.expectedPatchFilename)
+				}
+				assert.Equal(t, string(expected), afterPatch)
+			}
+		})
+	}
+}
diff --git a/pkg/git/patch_parser.go b/pkg/git/patch_parser.go
new file mode 100644
index 000000000..1dbacd01c
--- /dev/null
+++ b/pkg/git/patch_parser.go
@@ -0,0 +1,36 @@
+package git
+
+import (
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+type PatchParser struct {
+	Log *logrus.Entry
+}
+
+// NewPatchParser builds a new branch list builder
+func NewPatchParser(log *logrus.Entry) (*PatchParser, error) {
+	return &PatchParser{
+		Log: log,
+	}, nil
+}
+
+func (p *PatchParser) ParsePatch(patch string) ([]int, []int, error) {
+	lines := strings.Split(patch, "\n")
+	hunkStarts := []int{}
+	stageableLines := []int{}
+	pastHeader := false
+	for index, line := range lines {
+		if strings.HasPrefix(line, "@@") {
+			pastHeader = true
+			hunkStarts = append(hunkStarts, index)
+		}
+		if pastHeader && (strings.HasPrefix(line, "-") || strings.HasPrefix(line, "+")) {
+			stageableLines = append(stageableLines, index)
+		}
+	}
+	p.Log.WithField("staging", "staging").Info(stageableLines)
+	return hunkStarts, stageableLines, nil
+}
diff --git a/pkg/git/patch_parser_test.go b/pkg/git/patch_parser_test.go
new file mode 100644
index 000000000..6670aaea2
--- /dev/null
+++ b/pkg/git/patch_parser_test.go
@@ -0,0 +1,65 @@
+package git
+
+import (
+	"io/ioutil"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func newDummyPatchParser() *PatchParser {
+	return &PatchParser{
+		Log: newDummyLog(),
+	}
+}
+func TestParsePatch(t *testing.T) {
+	type scenario struct {
+		testName               string
+		patchFilename          string
+		shouldError            bool
+		expectedStageableLines []int
+		expectedHunkStarts     []int
+	}
+
+	scenarios := []scenario{
+		{
+			"Diff with one hunk",
+			"testdata/testPatchBefore.diff",
+			false,
+			[]int{8, 9, 10, 11},
+			[]int{4},
+		},
+		{
+			"Diff with two hunks",
+			"testdata/testPatchBefore2.diff",
+			false,
+			[]int{8, 9, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28, 33, 34, 35, 36, 37, 45, 46, 47, 48, 49, 50, 51, 52, 53},
+			[]int{4, 41},
+		},
+		{
+			"Unstaged file",
+			"testdata/addedFile.diff",
+			false,
+			[]int{6},
+			[]int{5},
+		},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			p := newDummyPatchParser()
+			beforePatch, err := ioutil.ReadFile(s.patchFilename)
+			if err != nil {
+				panic("Cannot open file at " + s.patchFilename)
+			}
+			hunkStarts, stageableLines, err := p.ParsePatch(string(beforePatch))
+			if s.shouldError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, s.expectedStageableLines, stageableLines)
+				assert.Equal(t, s.expectedHunkStarts, hunkStarts)
+			}
+		})
+	}
+}
diff --git a/pkg/git/testdata/addedFile.diff b/pkg/git/testdata/addedFile.diff
new file mode 100644
index 000000000..53966c4a1
--- /dev/null
+++ b/pkg/git/testdata/addedFile.diff
@@ -0,0 +1,7 @@
+diff --git a/blah b/blah
+new file mode 100644
+index 0000000..907b308
+--- /dev/null
++++ b/blah
+@@ -0,0 +1 @@
++blah
diff --git a/pkg/git/testdata/testPatchAfter1.diff b/pkg/git/testdata/testPatchAfter1.diff
new file mode 100644
index 000000000..88066e1c2
--- /dev/null
+++ b/pkg/git/testdata/testPatchAfter1.diff
@@ -0,0 +1,13 @@
+diff --git a/pkg/git/branch_list_builder.go b/pkg/git/branch_list_builder.go
+index 60ec4e0..db4485d 100644
+--- a/pkg/git/branch_list_builder.go
++++ b/pkg/git/branch_list_builder.go
+@@ -14,8 +14,7 @@ import (
+ 
+ // context:
+ // we want to only show 'safe' branches (ones that haven't e.g. been deleted)
+-// which `git branch -a` gives us, but we also want the recency data that
+ // git reflog gives us.
+ // So we get the HEAD, then append get the reflog branches that intersect with
+ // our safe branches, then add the remaining safe branches, ensuring uniqueness
+ // along the way
diff --git a/pkg/git/testdata/testPatchAfter2.diff b/pkg/git/testdata/testPatchAfter2.diff
new file mode 100644
index 000000000..0a17c2b67
--- /dev/null
+++ b/pkg/git/testdata/testPatchAfter2.diff
@@ -0,0 +1,14 @@
+diff --git a/pkg/git/branch_list_builder.go b/pkg/git/branch_list_builder.go
+index 60ec4e0..db4485d 100644
+--- a/pkg/git/branch_list_builder.go
++++ b/pkg/git/branch_list_builder.go
+@@ -14,8 +14,9 @@ import (
+ 
+ // context:
+ // we want to only show 'safe' branches (ones that haven't e.g. been deleted)
+ // which `git branch -a` gives us, but we also want the recency data that
+ // git reflog gives us.
++// test 2 - if I remove this, I decrement the end counter
+ // So we get the HEAD, then append get the reflog branches that intersect with
+ // our safe branches, then add the remaining safe branches, ensuring uniqueness
+ // along the way
diff --git a/pkg/git/testdata/testPatchAfter3.diff b/pkg/git/testdata/testPatchAfter3.diff
new file mode 100644
index 000000000..03492450d
--- /dev/null
+++ b/pkg/git/testdata/testPatchAfter3.diff
@@ -0,0 +1,25 @@
+diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
+index a8fc600..6d8f7d7 100644
+--- a/pkg/git/patch_modifier.go
++++ b/pkg/git/patch_modifier.go
+@@ -36,18 +36,19 @@ func (p *PatchModifier) ModifyPatchForHunk(patch string, hunkStarts []int, curre
+ 		hunkEnd = hunkStarts[nextHunkStartIndex]
+ 	}
+ 
+ 	headerLength := 4
+ 	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+ 	output += strings.Join(lines[hunkStart:hunkEnd], "\n") + "\n"
+ 
+ 	return output, nil
+ }
+ 
++func getHeaderLength(patchLines []string) (int, error) {
+ // ModifyPatchForLine takes the original patch, which may contain several hunks,
+ // and the line number of the line we want to stage
+ func (p *PatchModifier) ModifyPatchForLine(patch string, lineNumber int) (string, error) {
+ 	lines := strings.Split(patch, "\n")
+ 	headerLength := 4
+ 	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+ 
+ 	hunkStart, err := p.getHunkStart(lines, lineNumber)
+
diff --git a/pkg/git/testdata/testPatchAfter4.diff b/pkg/git/testdata/testPatchAfter4.diff
new file mode 100644
index 000000000..99f894d9d
--- /dev/null
+++ b/pkg/git/testdata/testPatchAfter4.diff
@@ -0,0 +1,19 @@
+diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
+index a8fc600..6d8f7d7 100644
+--- a/pkg/git/patch_modifier.go
++++ b/pkg/git/patch_modifier.go
+@@ -124,13 +140,14 @@ func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, line
+ // @@ -14,8 +14,9 @@ import (
+ func (p *PatchModifier) updatedHeader(currentHeader string, lineChanges int) (string, error) {
+ 	// current counter is the number after the second comma
+ 	re := regexp.MustCompile(`^[^,]+,[^,]+,(\d+)`)
+ 	matches := re.FindStringSubmatch(currentHeader)
+ 	if len(matches) < 2 {
+ 		re = regexp.MustCompile(`^[^,]+,[^+]+\+(\d+)`)
+ 		matches = re.FindStringSubmatch(currentHeader)
+ 	}
+ 	prevLengthString := matches[1]
++	prevLengthString := re.FindStringSubmatch(currentHeader)[1]
+ 
+ 	prevLength, err := strconv.Atoi(prevLengthString)
+ 	if err != nil {
diff --git a/pkg/git/testdata/testPatchBefore.diff b/pkg/git/testdata/testPatchBefore.diff
new file mode 100644
index 000000000..14e4b0e23
--- /dev/null
+++ b/pkg/git/testdata/testPatchBefore.diff
@@ -0,0 +1,15 @@
+diff --git a/pkg/git/branch_list_builder.go b/pkg/git/branch_list_builder.go
+index 60ec4e0..db4485d 100644
+--- a/pkg/git/branch_list_builder.go
++++ b/pkg/git/branch_list_builder.go
+@@ -14,8 +14,8 @@ import (
+ 
+ // context:
+ // we want to only show 'safe' branches (ones that haven't e.g. been deleted)
+-// which `git branch -a` gives us, but we also want the recency data that
+-// git reflog gives us.
++// test 2 - if I remove this, I decrement the end counter
++// test
+ // So we get the HEAD, then append get the reflog branches that intersect with
+ // our safe branches, then add the remaining safe branches, ensuring uniqueness
+ // along the way
diff --git a/pkg/git/testdata/testPatchBefore2.diff b/pkg/git/testdata/testPatchBefore2.diff
new file mode 100644
index 000000000..552c04f5e
--- /dev/null
+++ b/pkg/git/testdata/testPatchBefore2.diff
@@ -0,0 +1,57 @@
+diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
+index a8fc600..6d8f7d7 100644
+--- a/pkg/git/patch_modifier.go
++++ b/pkg/git/patch_modifier.go
+@@ -36,18 +36,34 @@ func (p *PatchModifier) ModifyPatchForHunk(patch string, hunkStarts []int, curre
+ 		hunkEnd = hunkStarts[nextHunkStartIndex]
+ 	}
+ 
+-	headerLength := 4
++	headerLength, err := getHeaderLength(lines)
++	if err != nil {
++		return "", err
++	}
++
+ 	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+ 	output += strings.Join(lines[hunkStart:hunkEnd], "\n") + "\n"
+ 
+ 	return output, nil
+ }
+ 
++func getHeaderLength(patchLines []string) (int, error) {
++	for index, line := range patchLines {
++		if strings.HasPrefix(line, "@@") {
++			return index, nil
++		}
++	}
++	return 0, errors.New("Could not find any hunks in this patch")
++}
++
+ // ModifyPatchForLine takes the original patch, which may contain several hunks,
+ // and the line number of the line we want to stage
+ func (p *PatchModifier) ModifyPatchForLine(patch string, lineNumber int) (string, error) {
+ 	lines := strings.Split(patch, "\n")
+-	headerLength := 4
++	headerLength, err := getHeaderLength(lines)
++	if err != nil {
++		return "", err
++	}
+ 	output := strings.Join(lines[0:headerLength], "\n") + "\n"
+ 
+ 	hunkStart, err := p.getHunkStart(lines, lineNumber)
+@@ -124,13 +140,8 @@ func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, line
+ // @@ -14,8 +14,9 @@ import (
+ func (p *PatchModifier) updatedHeader(currentHeader string, lineChanges int) (string, error) {
+ 	// current counter is the number after the second comma
+-	re := regexp.MustCompile(`^[^,]+,[^,]+,(\d+)`)
+-	matches := re.FindStringSubmatch(currentHeader)
+-	if len(matches) < 2 {
+-		re = regexp.MustCompile(`^[^,]+,[^+]+\+(\d+)`)
+-		matches = re.FindStringSubmatch(currentHeader)
+-	}
+-	prevLengthString := matches[1]
++	re := regexp.MustCompile(`(\d+) @@`)
++	prevLengthString := re.FindStringSubmatch(currentHeader)[1]