MultiFileDiffReader returns trailing content

diff/diff_test.go

@@ -2,6 +2,7 @@ package diff
 
 import (
 	"bytes"
+	"io"
 	"io/ioutil"
 	"path/filepath"
 	"reflect"
@@ -818,6 +819,52 @@ func TestParseMultiFileDiffAndPrintMultiFileDiff(t *testing.T) {
 	}
 }
 
+func TestParseMultiFileDiffAndPrintMultiFileDiffIncludingTrailingContent(t *testing.T) {
+	testInput, err := ioutil.ReadFile(filepath.Join("testdata", "sample_multi_file_trailing_content.diff"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedDiffs, err := ioutil.ReadFile(filepath.Join("testdata", "sample_multi_file_trailing_content_diffsonly.diff"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	diffReader := NewMultiFileDiffReader(bytes.NewReader(testInput))
+	var diffs []*FileDiff
+	trailingContent := ""
+	for {
+		var fd *FileDiff
+		var err error
+		fd, trailingContent, err = diffReader.ReadFileWithTrailingContent()
+		if fd != nil {
+			diffs = append(diffs, fd)
+		}
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Error(err)
+		}
+	}
+
+	if len(diffs) != 2 {
+		t.Errorf("expected 2 diffs, got %d", len(diffs))
+	}
+
+	printed, err := PrintMultiFileDiff(diffs)
+	if err != nil {
+		t.Errorf("PrintMultiFileDiff: %s", err)
+	}
+	if !bytes.Equal(printed, expectedDiffs) {
+		t.Errorf("printed multi-file diff != original multi-file diff\n\n# PrintMultiFileDiff output - Original:\n%s", cmp.Diff(expectedDiffs, printed))
+	}
+
+	expectedTrailingContent := "some trailing content"
+	if trailingContent != expectedTrailingContent {
+		t.Errorf("expected trailing content %s, got %s", expectedTrailingContent, trailingContent)
+	}
+}
+
 func TestNoNewlineAtEnd(t *testing.T) {
 	diffs := map[string]struct {
 		diff              string

diff/parse.go

@@ -46,6 +46,14 @@ type MultiFileDiffReader struct {
 // all hunks) from r. If there are no more files in the diff, it
 // returns error io.EOF.
 func (r *MultiFileDiffReader) ReadFile() (*FileDiff, error) {
+	fd, _, err := r.ReadFileWithTrailingContent()
+	return fd, err
+}
+
+// ReadFileWithTrailingContent reads the next file unified diff (including
+// headers and all hunks) from r, also returning any trailing content. If there
+// are no more files in the diff, it returns error io.EOF.
+func (r *MultiFileDiffReader) ReadFileWithTrailingContent() (*FileDiff, string, error) {
 	fr := &FileDiffReader{
 		line:           r.line,
 		offset:         r.offset,
@@ -59,23 +67,33 @@ func (r *MultiFileDiffReader) ReadFile() (*FileDiff, error) {
 		switch e := err.(type) {
 		case *ParseError:
 			if e.Err == ErrNoFileHeader || e.Err == ErrExtendedHeadersEOF {
-				return nil, io.EOF
+				// Any non-diff content preceding a valid diff is included in the
+				// extended headers of the following diff. In this way, mixed diff /
+				// non-diff content can be parsed. Trailing non-diff content is
+				// different: it doesn't make sense to return a FileDiff with only
+				// extended headers populated. Instead, we return any trailing content
+				// in case the caller needs it.
+				trailing := ""
+				if fd != nil {
+					trailing = strings.Join(fd.Extended, "\n")
+				}
+				return nil, trailing, io.EOF
 			}
-			return nil, err
+			return nil, "", err
 
 		case OverflowError:
 			r.nextFileFirstLine = []byte(e)
-			return fd, nil
+			return fd, "", nil
 
 		default:
-			return nil, err
+			return nil, "", err
 		}
 	}
 
 	// FileDiff is added/deleted file
 	// No further collection of hunks needed
 	if fd.NewName == "" {
-		return fd, nil
+		return fd, "", nil
 	}
 
 	// Before reading hunks, check to see if there are any. If there
@@ -87,7 +105,7 @@ func (r *MultiFileDiffReader) ReadFile() (*FileDiff, error) {
 	hr := fr.HunksReader()
 	line, err := r.reader.readLine()
 	if err != nil && err != io.EOF {
-		return fd, err
+		return fd, "", err
 	}
 	line = bytes.TrimSuffix(line, []byte{'\n'})
 	if bytes.HasPrefix(line, hunkPrefix) {
@@ -101,18 +119,18 @@ func (r *MultiFileDiffReader) ReadFile() (*FileDiff, error) {
 					// This just means we finished reading the hunks for the
 					// current file. See the ErrBadHunkLine doc for more info.
 					r.nextFileFirstLine = e.Line
-					return fd, nil
+					return fd, "", nil
 				}
 			}
-			return nil, err
+			return nil, "", err
 		}
 	} else {
 		// There weren't any hunks, so that line we peeked ahead at
 		// actually belongs to the next file. Put it back.
 		r.nextFileFirstLine = line
 	}
 
-	return fd, nil
+	return fd, "", nil
 }
 
 // ReadAllFiles reads all file unified diffs (including headers and all

diff/testdata/sample_multi_file_trailing_content.diff

@@ -0,0 +1,24 @@
+some leading content
+diff --git a/comment-last-line.sql b/comment-last-line.sql
+index 04a1655..97bd115 100644
+--- a/comment-last-line.sql
++++ b/comment-last-line.sql
+@@ -1,4 +1,4 @@
+ select 1;
++++ invalid SQL comment
+ select 2;
+ select 3;
+--- end of three queries
+some content between diffs
+diff --git a/query.sql b/query.sql
+index 9537d7b..234ef35 100644
+--- a/query.sql
++++ b/query.sql
+@@ -1,5 +1,4 @@
+ select 1;
+--- this is my query
+ select 2;
+ select 3;
+--- this is the last line
++++ invalid sql comment
+some trailing content

diff/testdata/sample_multi_file_trailing_content_diffsonly.diff

@@ -0,0 +1,23 @@
+some leading content
+diff --git a/comment-last-line.sql b/comment-last-line.sql
+index 04a1655..97bd115 100644
+--- a/comment-last-line.sql
++++ b/comment-last-line.sql
+@@ -1,4 +1,4 @@
+ select 1;
++++ invalid SQL comment
+ select 2;
+ select 3;
+--- end of three queries
+some content between diffs
+diff --git a/query.sql b/query.sql
+index 9537d7b..234ef35 100644
+--- a/query.sql
++++ b/query.sql
+@@ -1,5 +1,4 @@
+ select 1;
+--- this is my query
+ select 2;
+ select 3;
+--- this is the last line
++++ invalid sql comment