staging lines and hunks

8 files changed, 249 insertions, 17 deletions

diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
index 05afcf5ba..3c523232e 100644
--- a/pkg/git/patch_modifier.go
+++ b/pkg/git/patch_modifier.go
@@ -6,11 +6,14 @@ import (
 	"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
@@ -20,11 +23,49 @@ func NewPatchModifier(log *logrus.Entry) (*PatchModifier, error) {
 	}, nil
 }
 
-// ModifyPatch takes the original patch, which may contain several hunks,
+// 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) ModifyPatch(patch string, lineNumber int) (string, error) {
+func (p *PatchModifier) ModifyPatchForLine(patch string, lineNumber int) (string, error) {
 	lines := strings.Split(patch, "\n")
-	headerLength := 4
+	headerLength, err := p.getHeaderLength(lines)
+	if err != nil {
+		return "", err
+	}
 	output := strings.Join(lines[0:headerLength], "\n") + "\n"
 
 	hunkStart, err := p.getHunkStart(lines, lineNumber)
@@ -55,7 +96,8 @@ func (p *PatchModifier) getHunkStart(patchLines []string, lineNumber int) (int,
 			return hunkStart, nil
 		}
 	}
-	return 0, errors.New("Could not find hunk")
+
+	return 0, errors.New(p.Tr.SLocalize("CantFindHunk"))
 }
 
 func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, lineNumber int) ([]string, error) {
@@ -101,13 +143,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]
 
 	prevLength, err := strconv.Atoi(prevLengthString)
 	if err != nil {
diff --git a/pkg/git/patch_modifier_test.go b/pkg/git/patch_modifier_test.go
index af7be3751..bc2073d55 100644
--- a/pkg/git/patch_modifier_test.go
+++ b/pkg/git/patch_modifier_test.go
@@ -19,7 +19,7 @@ func newDummyPatchModifier() *PatchModifier {
 		Log: newDummyLog(),
 	}
 }
-func TestModifyPatch(t *testing.T) {
+func TestModifyPatchForLine(t *testing.T) {
 	type scenario struct {
 		testName              string
 		patchFilename         string
@@ -43,6 +43,27 @@ func TestModifyPatch(t *testing.T) {
 			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 {
@@ -52,7 +73,7 @@ func TestModifyPatch(t *testing.T) {
 			if err != nil {
 				panic("Cannot open file at " + s.patchFilename)
 			}
-			afterPatch, err := p.ModifyPatch(string(beforePatch), s.lineNumber)
+			afterPatch, err := p.ModifyPatchForLine(string(beforePatch), s.lineNumber)
 			if s.shouldError {
 				assert.Error(t, err)
 			} else {
diff --git a/pkg/git/patch_parser.go b/pkg/git/patch_parser.go
index 8fa4d355b..1dbacd01c 100644
--- a/pkg/git/patch_parser.go
+++ b/pkg/git/patch_parser.go
@@ -21,15 +21,16 @@ func (p *PatchParser) ParsePatch(patch string) ([]int, []int, error) {
 	lines := strings.Split(patch, "\n")
 	hunkStarts := []int{}
 	stageableLines := []int{}
-	headerLength := 4
-	for offsetIndex, line := range lines[headerLength:] {
-		index := offsetIndex + headerLength
+	pastHeader := false
+	for index, line := range lines {
 		if strings.HasPrefix(line, "@@") {
+			pastHeader = true
 			hunkStarts = append(hunkStarts, index)
 		}
-		if strings.HasPrefix(line, "-") || strings.HasPrefix(line, "+") {
+		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/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/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]
+ 
+ 	prevLength, err := strconv.Atoi(prevLengthString)
+ 	if err != nil {