github.com/dawidd6/lazygit@v0.8.1/pkg/git/patch_modifier.go (about) 1 package git 2 3 import ( 4 "regexp" 5 "strconv" 6 "strings" 7 8 "github.com/go-errors/errors" 9 10 "github.com/jesseduffield/lazygit/pkg/i18n" 11 "github.com/jesseduffield/lazygit/pkg/utils" 12 "github.com/sirupsen/logrus" 13 ) 14 15 type PatchModifier struct { 16 Log *logrus.Entry 17 Tr *i18n.Localizer 18 } 19 20 // NewPatchModifier builds a new branch list builder 21 func NewPatchModifier(log *logrus.Entry) (*PatchModifier, error) { 22 return &PatchModifier{ 23 Log: log, 24 }, nil 25 } 26 27 // ModifyPatchForHunk takes the original patch, which may contain several hunks, 28 // and removes any hunks that aren't the selected hunk 29 func (p *PatchModifier) ModifyPatchForHunk(patch string, hunkStarts []int, currentLine int) (string, error) { 30 // get hunk start and end 31 lines := strings.Split(patch, "\n") 32 hunkStartIndex := utils.PrevIndex(hunkStarts, currentLine) 33 hunkStart := hunkStarts[hunkStartIndex] 34 nextHunkStartIndex := utils.NextIndex(hunkStarts, currentLine) 35 var hunkEnd int 36 if nextHunkStartIndex == 0 { 37 hunkEnd = len(lines) - 1 38 } else { 39 hunkEnd = hunkStarts[nextHunkStartIndex] 40 } 41 42 headerLength, err := p.getHeaderLength(lines) 43 if err != nil { 44 return "", err 45 } 46 47 output := strings.Join(lines[0:headerLength], "\n") + "\n" 48 output += strings.Join(lines[hunkStart:hunkEnd], "\n") + "\n" 49 50 return output, nil 51 } 52 53 func (p *PatchModifier) getHeaderLength(patchLines []string) (int, error) { 54 for index, line := range patchLines { 55 if strings.HasPrefix(line, "@@") { 56 return index, nil 57 } 58 } 59 return 0, errors.New(p.Tr.SLocalize("CantFindHunks")) 60 } 61 62 // ModifyPatchForLine takes the original patch, which may contain several hunks, 63 // and the line number of the line we want to stage 64 func (p *PatchModifier) ModifyPatchForLine(patch string, lineNumber int) (string, error) { 65 lines := strings.Split(patch, "\n") 66 headerLength, err := p.getHeaderLength(lines) 67 if err != nil { 68 return "", err 69 } 70 output := strings.Join(lines[0:headerLength], "\n") + "\n" 71 72 hunkStart, err := p.getHunkStart(lines, lineNumber) 73 if err != nil { 74 return "", err 75 } 76 77 hunk, err := p.getModifiedHunk(lines, hunkStart, lineNumber) 78 if err != nil { 79 return "", err 80 } 81 82 output += strings.Join(hunk, "\n") 83 84 return output, nil 85 } 86 87 // getHunkStart returns the line number of the hunk we're going to be modifying 88 // in order to stage our line 89 func (p *PatchModifier) getHunkStart(patchLines []string, lineNumber int) (int, error) { 90 // find the hunk that we're modifying 91 hunkStart := 0 92 for index, line := range patchLines { 93 if strings.HasPrefix(line, "@@") { 94 hunkStart = index 95 } 96 if index == lineNumber { 97 return hunkStart, nil 98 } 99 } 100 101 return 0, errors.New(p.Tr.SLocalize("CantFindHunk")) 102 } 103 104 func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, lineNumber int) ([]string, error) { 105 lineChanges := 0 106 // strip the hunk down to just the line we want to stage 107 newHunk := []string{patchLines[hunkStart]} 108 for offsetIndex, line := range patchLines[hunkStart+1:] { 109 index := offsetIndex + hunkStart + 1 110 if strings.HasPrefix(line, "@@") { 111 newHunk = append(newHunk, "\n") 112 break 113 } 114 if index != lineNumber { 115 // we include other removals but treat them like context 116 if strings.HasPrefix(line, "-") { 117 newHunk = append(newHunk, " "+line[1:]) 118 lineChanges += 1 119 continue 120 } 121 // we don't include other additions 122 if strings.HasPrefix(line, "+") { 123 lineChanges -= 1 124 continue 125 } 126 } 127 newHunk = append(newHunk, line) 128 } 129 130 var err error 131 newHunk[0], err = p.updatedHeader(newHunk[0], lineChanges) 132 if err != nil { 133 return nil, err 134 } 135 136 return newHunk, nil 137 } 138 139 // updatedHeader returns the hunk header with the updated line range 140 // we need to update the hunk length to reflect the changes we made 141 // if the hunk has three additions but we're only staging one, then 142 // @@ -14,8 +14,11 @@ import ( 143 // becomes 144 // @@ -14,8 +14,9 @@ import ( 145 func (p *PatchModifier) updatedHeader(currentHeader string, lineChanges int) (string, error) { 146 // current counter is the number after the second comma 147 re := regexp.MustCompile(`(\d+) @@`) 148 prevLengthString := re.FindStringSubmatch(currentHeader)[1] 149 150 prevLength, err := strconv.Atoi(prevLengthString) 151 if err != nil { 152 return "", err 153 } 154 re = regexp.MustCompile(`\d+ @@`) 155 newLength := strconv.Itoa(prevLength + lineChanges) 156 return re.ReplaceAllString(currentHeader, newLength+" @@"), nil 157 }