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  }