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