From 4ebce9caa45902d8ca5f63a242d53e1e5c6072d3 Mon Sep 17 00:00:00 2001
From: michaelcheah <20640150+michaelcheah@users.noreply.github.com>
Date: Fri, 24 Mar 2023 19:37:54 +0000
Subject: [PATCH] fix: colored display for diffs with newlines

---
 diffmatchpatch/diff.go      | 27 +++++++++++++++++++++------
 diffmatchpatch/diff_test.go | 11 +++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/diffmatchpatch/diff.go b/diffmatchpatch/diff.go
index 4f7b424..a312c1a 100644
--- a/diffmatchpatch/diff.go
+++ b/diffmatchpatch/diff.go
@@ -1151,13 +1151,28 @@ func (dmp *DiffMatchPatch) DiffPrettyText(diffs []Diff) string {
 
 		switch diff.Type {
 		case DiffInsert:
-			_, _ = buff.WriteString("\x1b[32m")
-			_, _ = buff.WriteString(text)
-			_, _ = buff.WriteString("\x1b[0m")
+			lines := strings.Split(text, "\n")
+			for i, line := range lines {
+				_, _ = buff.WriteString("\x1b[32m")
+				_, _ = buff.WriteString(line)
+				if i < len(lines)-1 {
+					_, _ = buff.WriteString("\x1b[0m\n")
+				} else {
+					_, _ = buff.WriteString("\x1b[0m")
+				}
+			}
+
 		case DiffDelete:
-			_, _ = buff.WriteString("\x1b[31m")
-			_, _ = buff.WriteString(text)
-			_, _ = buff.WriteString("\x1b[0m")
+			lines := strings.Split(text, "\n")
+			for i, line := range lines {
+				_, _ = buff.WriteString("\x1b[31m")
+				_, _ = buff.WriteString(line)
+				if i < len(lines)-1 {
+					_, _ = buff.WriteString("\x1b[0m\n")
+				} else {
+					_, _ = buff.WriteString("\x1b[0m")
+				}
+			}
 		case DiffEqual:
 			_, _ = buff.WriteString(text)
 		}
diff --git a/diffmatchpatch/diff_test.go b/diffmatchpatch/diff_test.go
index d6fed50..04329d6 100644
--- a/diffmatchpatch/diff_test.go
+++ b/diffmatchpatch/diff_test.go
@@ -1033,6 +1033,17 @@ func TestDiffPrettyText(t *testing.T) {
 
 			Expected: "a\n\x1b[31m<B>b</B>\x1b[0m\x1b[32mc&d\x1b[0m",
 		},
+		{
+			Diffs: []Diff{
+				{Type: DiffEqual, Text: "a\n"},
+				{Type: DiffDelete, Text: "b\nc\n"},
+				{Type: DiffEqual, Text: "def"},
+				{Type: DiffInsert, Text: "\ng\nh"},
+				{Type: DiffEqual, Text: "\ni"},
+			},
+
+			Expected: "a\n\x1b[31mb\x1b[0m\n\x1b[31mc\x1b[0m\n\x1b[31m\x1b[0mdef\x1b[32m\x1b[0m\n\x1b[32mg\x1b[0m\n\x1b[32mh\x1b[0m\ni",
+		},
 	} {
 		actual := dmp.DiffPrettyText(tc.Diffs)
 		assert.Equal(t, tc.Expected, actual, fmt.Sprintf("Test case #%d, %#v", i, tc))