Handle no-newline-at-end-of-file for orig

Author: sqs

diff/diff.go

@@ -36,8 +36,9 @@ func (d *FileDiff) Stat() Stat {
 // A Hunk represents a series of changes (additions or deletions) in a
 // file's unified diff.
 type Hunk struct {
-	OrigStartLine int // starting line number in original file
-	OrigLines     int // number of lines the hunk applies to in the original file
+	OrigStartLine   int // starting line number in original file
+	OrigLines       int // number of lines the hunk applies to in the original file
+	OrigNoNewlineAt int // if > 0, then the original file had a 'No newline at end of file' mark at this offset
 
 	NewStartLine int // starting line number in new file
 	NewLines     int // number of lines the hunk applies to in the new file

diff/diff_test.go

@@ -54,6 +54,13 @@ func TestParseHunksAndPrintHunks(t *testing.T) {
 		{
 			filename: "sample_hunks_no_newline.diff",
 		},
+		{filename: "no_newline_both.diff"},
+		{filename: "no_newline_both2.diff"},
+		{filename: "no_newline_orig.diff"},
+		{filename: "no_newline_new.diff"},
+		{filename: "empty_orig.diff"},
+		{filename: "empty_new.diff"},
+		{filename: "oneline_hunk.diff"},
 		{filename: "empty.diff"},
 	}
 	for _, test := range tests {
@@ -162,14 +169,14 @@ func TestParseMultiFileDiffAndPrintMultiFileDiff(t *testing.T) {
 
 func TestNoNewlineAtEnd(t *testing.T) {
 	orig := `@@ -1,1 +1,1 @@
--b
+-a
 +b
 \ No newline at end of file
 `
 
 	hunks, err := ParseHunks([]byte(orig))
 	if err != nil {
-		t.Fatal("ParseHunks: %s", err)
+		t.Fatalf("ParseHunks: %s", err)
 	}
 
 	for _, hunk := range hunks {
@@ -182,7 +189,7 @@ func TestNoNewlineAtEnd(t *testing.T) {
 
 		printed, err := PrintHunks(hunks)
 		if err != nil {
-			t.Fatal("PrintHunks: %s", err)
+			t.Fatalf("PrintHunks: %s", err)
 		}
 		if printed := string(printed); printed != orig {
 			t.Errorf("printed diff hunks != original diff hunks\n\n# PrintHunks output:\n%q\n\n# Original:\n%q", printed, orig)

diff/parse.go

@@ -324,6 +324,7 @@ type HunksReader struct {
 // returns error io.EOF.
 func (r *HunksReader) ReadHunk() (*Hunk, error) {
 	r.hunk = nil
+	lastLineFromOrig := true
 	var line []byte
 	for {
 		if r.nextHunkHeaderLine != nil {
@@ -392,13 +393,26 @@ func (r *HunksReader) ReadHunk() (*Hunk, error) {
 				return r.hunk, &ParseError{r.line, r.offset, &ErrBadHunkLine{Line: line}}
 			}
 			if bytes.Equal(line, []byte(noNewlineMessage)) {
-				// Remove previous line's newline.
-				if len(r.hunk.Body) != 0 {
-					r.hunk.Body = r.hunk.Body[:len(r.hunk.Body)-1]
+				if lastLineFromOrig {
+					// Retain the newline in the body (otherwise the
+					// diff line would be like "-a+b", where "+b" is
+					// the the next line of the new file, which is not
+					// validly formatted) but record that the orig had
+					// no newline.
+					r.hunk.OrigNoNewlineAt = len(r.hunk.Body)
+				} else {
+					// Remove previous line's newline.
+					if len(r.hunk.Body) != 0 {
+						r.hunk.Body = r.hunk.Body[:len(r.hunk.Body)-1]
+					}
 				}
 				continue
 			}
 
+			if len(line) > 0 {
+				lastLineFromOrig = line[0] == '-'
+			}
+
 			r.hunk.Body = append(r.hunk.Body, line...)
 			r.hunk.Body = append(r.hunk.Body, '\n')
 		}

diff/print.go

@@ -86,20 +86,41 @@ func PrintHunks(hunks []*Hunk) ([]byte, error) {
 		if _, err := fmt.Fprintln(&buf); err != nil {
 			return nil, err
 		}
-		if _, err := buf.Write(hunk.Body); err != nil {
-			return nil, err
-		}
-		if !bytes.HasSuffix(hunk.Body, []byte{'\n'}) {
-			if _, err := fmt.Fprintln(&buf); err != nil {
+
+		if hunk.OrigNoNewlineAt == 0 {
+			if _, err := buf.Write(hunk.Body); err != nil {
 				return nil, err
 			}
-			if _, err := buf.Write([]byte(noNewlineMessage)); err != nil {
+		} else {
+			if _, err := buf.Write(hunk.Body[:hunk.OrigNoNewlineAt]); err != nil {
 				return nil, err
 			}
+			if err := printNoNewlineMessage(&buf); err != nil {
+				return nil, err
+			}
+			if _, err := buf.Write(hunk.Body[hunk.OrigNoNewlineAt:]); err != nil {
+				return nil, err
+			}
+		}
+
+		if !bytes.HasSuffix(hunk.Body, []byte{'\n'}) {
 			if _, err := fmt.Fprintln(&buf); err != nil {
 				return nil, err
 			}
+			if err := printNoNewlineMessage(&buf); err != nil {
+				return nil, err
+			}
 		}
 	}
 	return buf.Bytes(), nil
 }
+
+func printNoNewlineMessage(w io.Writer) error {
+	if _, err := w.Write([]byte(noNewlineMessage)); err != nil {
+		return err
+	}
+	if _, err := fmt.Fprintln(w); err != nil {
+		return err
+	}
+	return nil
+}

diff/testdata/empty_new.diff

@@ -0,0 +1,2 @@
+@@ -1,1 +0,0 @@
+-b

diff/testdata/empty_orig.diff

@@ -0,0 +1,2 @@
+@@ -0,0 +1,1 @@
++b

diff/testdata/no_newline_both.diff

@@ -0,0 +1,5 @@
+@@ -1,1 +1,1 @@
+-a
+\ No newline at end of file
++b
+\ No newline at end of file

diff/testdata/no_newline_both2.diff

@@ -0,0 +1,8 @@
+@@ -1,3 +1,3 @@
+ 0
+-a
+-a
+\ No newline at end of file
++b
++b
+\ No newline at end of file

diff/testdata/no_newline_new.diff

@@ -0,0 +1,6 @@
+@@ -1,3 +1,2 @@
+ a
+-a
+-a
++a
+\ No newline at end of file

diff/testdata/no_newline_orig.diff

@@ -0,0 +1,4 @@
+@@ -1,1 +1,1 @@
+-a
+\ No newline at end of file
++b

diff/testdata/oneline_hunk.diff

@@ -0,0 +1,3 @@
+@@ -1,1 +1,1 @@
+-a
++b