initial support for staging individual lines

2 files changed, 48 insertions, 4 deletions

diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
index 0e31af952..05afcf5ba 100644
--- a/pkg/git/patch_modifier.go
+++ b/pkg/git/patch_modifier.go
@@ -61,9 +61,13 @@ func (p *PatchModifier) getHunkStart(patchLines []string, lineNumber int) (int,
 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{}
-	for offsetIndex, line := range patchLines[hunkStart:] {
-		index := offsetIndex + hunkStart
+	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, "-") {
@@ -98,7 +102,12 @@ func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, line
 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]
+	matches := re.FindStringSubmatch(currentHeader)
+	if len(matches) < 2 {
+		re = regexp.MustCompile(`^[^,]+,[^+]+\+(\d+)`)
+		matches = re.FindStringSubmatch(currentHeader)
+	}
+	prevLengthString := matches[1]
 
 	prevLength, err := strconv.Atoi(prevLengthString)
 	if err != nil {
diff --git a/pkg/git/patch_parser.go b/pkg/git/patch_parser.go
new file mode 100644
index 000000000..8fa4d355b
--- /dev/null
+++ b/pkg/git/patch_parser.go
@@ -0,0 +1,35 @@
+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{}
+	headerLength := 4
+	for offsetIndex, line := range lines[headerLength:] {
+		index := offsetIndex + headerLength
+		if strings.HasPrefix(line, "@@") {
+			hunkStarts = append(hunkStarts, index)
+		}
+		if strings.HasPrefix(line, "-") || strings.HasPrefix(line, "+") {
+			stageableLines = append(stageableLines, index)
+		}
+	}
+	return hunkStarts, stageableLines, nil
+}