Remove color escape after newline from diff line (#2)

* remove color escape after newline from diff line

* clear color codes after emitting format line, tests for color codes

Author: whwalter

difflib/difflib.go

@@ -595,7 +595,7 @@ func (d *Differ) StructuredDump(tag byte, x []string, low int, high int) (out []
 	size := high - low
 	out = make([]DiffLine, size)
 	for i := 0; i < size; i++ {
-		out[i] = NewDiffLine(tag, x[i + low])
+		out[i] = NewDiffLine(tag, x[i+low])
 	}
 	return out
 }
@@ -877,9 +877,8 @@ func WriteUnifiedDiff(writer io.Writer, diff LineDiffParams) error {
 		equalFormat := " %s"
 
 		if diff.Colored {
-			addFormat = "\x1b[32m+%s"    // Foreground Green
-			removeFormat = "\x1b[31m-%s" // Foreground Red
-			equalFormat = "\x1b[0m %s"   // No Color
+			addFormat = "\x1b[32m+%s\x1b[0m"    // Foreground Green
+			removeFormat = "\x1b[31m-%s\x1b[0m" // Foreground Red
 		}
 
 		for _, c := range g {

difflib/difflib_test.go

@@ -615,3 +615,46 @@ func TestGetUnifiedDiffString(t *testing.T) {
 		t.Fatal("GetUnifiedDiffString failure:", diffStr)
 	}
 }
+
+func TestColoredDiffString(t *testing.T) {
+	A := "one\ntwo\nthree\nfour"
+	B := "one\ntwo\nthr33\nfour"
+
+	// Build diff
+	diff := LineDiffParams{
+		A:       SplitLines(A),
+		B:       SplitLines(B),
+		Context: 3,
+		Colored: true,
+	}
+	// string builder always returns a nil error so it's safe to ignore here
+	diffStr, err := GetUnifiedDiffString(diff)
+	if err != nil {
+		t.Fatal("Failed")
+	}
+
+	if !strings.Contains(diffStr, "\x1b[31m-three\n\x1b[0m") {
+		t.Fatal("Failed to remove removeFormat formatting")
+	}
+
+	if !strings.Contains(diffStr, "one\n") {
+		t.Fatal("Neutal Line contained color format")
+	}
+
+	// Reverse diff
+	diff = LineDiffParams{
+		A:       SplitLines(B),
+		B:       SplitLines(A),
+		Context: 3,
+		Colored: true,
+	}
+	// string builder always returns a nil error so it's safe to ignore here
+	diffStr, err = GetUnifiedDiffString(diff)
+	if err != nil {
+		t.Fatal("Failed")
+	}
+
+	if !strings.Contains(diffStr, "\x1b[32m+three\n\x1b[0m") {
+		t.Fatal("Failed to remove addFormat formatting")
+	}
+}