add patch modifier struct

5 files changed, 220 insertions, 0 deletions

diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
new file mode 100644
index 000000000..0e31af952
--- /dev/null
+++ b/pkg/git/patch_modifier.go
@@ -0,0 +1,110 @@
+package git
+
+import (
+	"errors"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+type PatchModifier struct {
+	Log *logrus.Entry
+}
+
+// NewPatchModifier builds a new branch list builder
+func NewPatchModifier(log *logrus.Entry) (*PatchModifier, error) {
+	return &PatchModifier{
+		Log: log,
+	}, nil
+}
+
+// ModifyPatch 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) {
+	lines := strings.Split(patch, "\n")
+	headerLength := 4
+	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("Could not find hunk")
+}
+
+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
+		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..af7be3751
--- /dev/null
+++ b/pkg/git/patch_modifier_test.go
@@ -0,0 +1,68 @@
+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 TestModifyPatch(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",
+		},
+	}
+
+	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.ModifyPatch(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/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/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