From 4e94726054bf6b698765555c45750f2207b0ac56 Mon Sep 17 00:00:00 2001
From: Matthias Diester <mddiester@gmail.com>
Date: Wed, 19 Jun 2019 23:05:49 +0200
Subject: [PATCH] Add line by line coloring

Add EachLine() style option, which enables a flag that coloring should not be
applied to new line sequences. Each lines has its own color sequences then.
---
 convenience.go      | 74 ++++++++++++++++++++++++++++++++-------------
 convenience_test.go | 14 +++++++++
 2 files changed, 67 insertions(+), 21 deletions(-)

diff --git a/convenience.go b/convenience.go
index 7ac576d..28b2f6f 100644
--- a/convenience.go
+++ b/convenience.go
@@ -28,8 +28,11 @@ import (
 	colorful "github.com/lucasb-eyer/go-colorful"
 )
 
-// StyleOption defines style option for strings
-type StyleOption func(*String)
+// StyleOption defines style option for colored strings
+type StyleOption struct {
+	flags       []string
+	postProcess func(*String, map[string]struct{})
+}
 
 // PlainTextLength returns the length of the input text without any escape
 // sequences.
@@ -66,39 +69,61 @@ func Substring(text string, start int, end int) string {
 
 // Bold applies the bold text parameter
 func Bold() StyleOption {
-	return func(s *String) {
-		for i := range *s {
-			(*s)[i].Settings |= 1 << 2
-		}
+	return StyleOption{
+		postProcess: func(s *String, flags map[string]struct{}) {
+			for i := range *s {
+				(*s)[i].Settings |= 1 << 2
+			}
+		},
 	}
 }
 
 // Italic applies the italic text parameter
 func Italic() StyleOption {
-	return func(s *String) {
-		for i := range *s {
-			(*s)[i].Settings |= 1 << 3
-		}
+	return StyleOption{
+		postProcess: func(s *String, flags map[string]struct{}) {
+			for i := range *s {
+				(*s)[i].Settings |= 1 << 3
+			}
+		},
 	}
 }
 
 // Foreground sets the given color as the foreground color of the text
 func Foreground(color colorful.Color) StyleOption {
-	r, g, b := color.RGB255()
-	return func(s *String) {
-		for i := range *s {
-			(*s)[i].Settings |= 1
-			(*s)[i].Settings |= uint64(r) << 8
-			(*s)[i].Settings |= uint64(g) << 16
-			(*s)[i].Settings |= uint64(b) << 24
-		}
+	return StyleOption{
+		postProcess: func(s *String, flags map[string]struct{}) {
+			r, g, b := color.RGB255()
+			_, skipNewLine := flags["skipNewLine"]
+
+			for i := range *s {
+				if skipNewLine && (*s)[i].Symbol == '\n' {
+					continue
+				}
+
+				(*s)[i].Settings |= 1
+				(*s)[i].Settings |= uint64(r) << 8
+				(*s)[i].Settings |= uint64(g) << 16
+				(*s)[i].Settings |= uint64(b) << 24
+			}
+		},
 	}
 }
 
 // EnableTextAnnotations enables post-processing to evaluate text annotations
 func EnableTextAnnotations() StyleOption {
-	return func(s *String) {
-		processTextAnnotations(s)
+	return StyleOption{
+		postProcess: func(s *String, flags map[string]struct{}) {
+			processTextAnnotations(s)
+		},
+	}
+}
+
+// EachLine enables that new line sequences will be ignored during coloring,
+// which will lead to strings that are colored line by line and not as a block.
+func EachLine() StyleOption {
+	return StyleOption{
+		flags: []string{"skipNewLine"},
 	}
 }
 
@@ -111,8 +136,15 @@ func Style(text string, styleOptions ...StyleOption) string {
 		panic(err)
 	}
 
+	flags := map[string]struct{}{}
 	for _, styleOption := range styleOptions {
-		styleOption(result)
+		for _, flag := range styleOption.flags {
+			flags[flag] = struct{}{}
+		}
+
+		if styleOption.postProcess != nil {
+			styleOption.postProcess(result, flags)
+		}
 	}
 
 	return result.String()
diff --git a/convenience_test.go b/convenience_test.go
index 5b96803..4787f77 100644
--- a/convenience_test.go
+++ b/convenience_test.go
@@ -90,5 +90,19 @@ var _ = Describe("convenience functions", func() {
 			Expect(Style("_text_", Foreground(YellowGreen), EnableTextAnnotations())).To(
 				BeEquivalentTo("\x1b[3;38;2;154;205;50mtext\x1b[0m"))
 		})
+
+		It("should support both line by line coloring as well as full block coloring", func() {
+			// By default, color the whole string including new line sequences
+			Expect(Style("text\ntext", Foreground(Yellow))).To(
+				BeEquivalentTo("\x1b[38;2;255;255;0mtext\ntext\x1b[0m"))
+
+			// If EachLine is enabled before coloring, ignore new line sequences
+			Expect(Style("text\ntext", EachLine(), Foreground(Yellow))).To(
+				BeEquivalentTo("\x1b[38;2;255;255;0mtext\x1b[0m\n\x1b[38;2;255;255;0mtext\x1b[0m"))
+
+			// If EachLine is enabled after coloring, it has no effect
+			Expect(Style("text\ntext", Foreground(Yellow), EachLine())).To(
+				BeEquivalentTo("\x1b[38;2;255;255;0mtext\ntext\x1b[0m"))
+		})
 	})
 })