text: Escape and StripEscape (#80)

text/ansi.go

@@ -1,5 +1,55 @@
 package text
 
+import "strings"
+
 // ANSICodesSupported will be true on consoles where ANSI Escape Codes/Sequences
 // are supported.
 var ANSICodesSupported = areANSICodesSupported()
+
+// Escape encodes the string with the ANSI Escape Sequence.
+// For ex.:
+//  Escape("Ghost", "") == "Ghost"
+//	Escape("Ghost", "\x1b[91m") == "\x1b[91mGhost\x1b[0m"
+//	Escape("\x1b[94mGhost\x1b[0mLady", "\x1b[91m") == "\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"
+//	Escape("Nymeria\x1b[94mGhost\x1b[0mLady", "\x1b[91m") == "\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"
+//	Escape("Nymeria \x1b[94mGhost\x1b[0m Lady", "\x1b[91m") == "\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m"
+func Escape(str string, escapeSeq string) string {
+	out := ""
+	if !strings.HasPrefix(str, EscapeStart) {
+		out += escapeSeq
+	}
+	out += strings.Replace(str, EscapeReset, EscapeReset+escapeSeq, -1)
+	if !strings.HasSuffix(out, EscapeReset) {
+		out += EscapeReset
+	}
+	if strings.Contains(out, escapeSeq+EscapeReset) {
+		out = strings.Replace(out, escapeSeq+EscapeReset, "", -1)
+	}
+	return out
+}
+
+// StripEscape strips all ANSI Escape Sequence from the string.
+// For ex.:
+//  StripEscape("Ghost") == "Ghost"
+//	StripEscape("\x1b[91mGhost\x1b[0m") == "Ghost"
+//	StripEscape("\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m") == "GhostLady"
+//	StripEscape("\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m") == "NymeriaGhostLady"
+//	StripEscape("\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m") == "Nymeria Ghost Lady"
+func StripEscape(str string) string {
+	var out strings.Builder
+	out.Grow(RuneCount(str))
+
+	isEscSeq := false
+	for _, sChr := range str {
+		if sChr == EscapeStartRune {
+			isEscSeq = true
+		}
+		if !isEscSeq {
+			out.WriteRune(sChr)
+		}
+		if isEscSeq && sChr == EscapeStopRune {
+			isEscSeq = false
+		}
+	}
+	return out.String()
+}

text/ansi_test.go

@@ -0,0 +1,50 @@
+package text
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEscape(t *testing.T) {
+	assert.Equal(t, "\x1b[91mGhost\x1b[0m", Escape("Ghost", FgHiRed.EscapeSeq()))
+	assert.Equal(t, "\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m", Escape(FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq()))
+	assert.Equal(t, "\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m", Escape("Nymeria"+FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq()))
+	assert.Equal(t, "\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m", Escape("Nymeria "+FgHiBlue.Sprint("Ghost")+" Lady", FgHiRed.EscapeSeq()))
+}
+
+func ExampleEscape() {
+	fmt.Printf("Escape(%#v, %#v) == %#v\n", "Ghost", "", Escape("Ghost", ""))
+	fmt.Printf("Escape(%#v, %#v) == %#v\n", "Ghost", FgHiRed.EscapeSeq(), Escape("Ghost", FgHiRed.EscapeSeq()))
+	fmt.Printf("Escape(%#v, %#v) == %#v\n", FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq(), Escape(FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq()))
+	fmt.Printf("Escape(%#v, %#v) == %#v\n", "Nymeria"+FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq(), Escape("Nymeria"+FgHiBlue.Sprint("Ghost")+"Lady", FgHiRed.EscapeSeq()))
+	fmt.Printf("Escape(%#v, %#v) == %#v\n", "Nymeria "+FgHiBlue.Sprint("Ghost")+" Lady", FgHiRed.EscapeSeq(), Escape("Nymeria "+FgHiBlue.Sprint("Ghost")+" Lady", FgHiRed.EscapeSeq()))
+
+	// Output: Escape("Ghost", "") == "Ghost"
+	// Escape("Ghost", "\x1b[91m") == "\x1b[91mGhost\x1b[0m"
+	// Escape("\x1b[94mGhost\x1b[0mLady", "\x1b[91m") == "\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"
+	// Escape("Nymeria\x1b[94mGhost\x1b[0mLady", "\x1b[91m") == "\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"
+	// Escape("Nymeria \x1b[94mGhost\x1b[0m Lady", "\x1b[91m") == "\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m"
+}
+
+func TestStripEscape(t *testing.T) {
+	assert.Equal(t, "Ghost", StripEscape(FgHiRed.Sprint("Ghost")))
+	assert.Equal(t, "GhostLady", StripEscape(FgHiBlue.Sprint("Ghost")+"Lady"))
+	assert.Equal(t, "NymeriaGhostLady", StripEscape("Nymeria"+FgHiBlue.Sprint("Ghost")+"Lady"))
+	assert.Equal(t, "Nymeria Ghost Lady", StripEscape("Nymeria "+FgHiBlue.Sprint("Ghost")+" Lady"))
+}
+
+func ExampleStripEscape() {
+	fmt.Printf("StripEscape(%#v) == %#v\n", "Ghost", StripEscape("Ghost"))
+	fmt.Printf("StripEscape(%#v) == %#v\n", "\x1b[91mGhost\x1b[0m", StripEscape("\x1b[91mGhost\x1b[0m"))
+	fmt.Printf("StripEscape(%#v) == %#v\n", "\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m", StripEscape("\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"))
+	fmt.Printf("StripEscape(%#v) == %#v\n", "\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m", StripEscape("\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m"))
+	fmt.Printf("StripEscape(%#v) == %#v\n", "\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m", StripEscape("\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m"))
+
+	// Output: StripEscape("Ghost") == "Ghost"
+	// StripEscape("\x1b[91mGhost\x1b[0m") == "Ghost"
+	// StripEscape("\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m") == "GhostLady"
+	// StripEscape("\x1b[91mNymeria\x1b[94mGhost\x1b[0m\x1b[91mLady\x1b[0m") == "NymeriaGhostLady"
+	// StripEscape("\x1b[91mNymeria \x1b[94mGhost\x1b[0m\x1b[91m Lady\x1b[0m") == "Nymeria Ghost Lady"
+}

text/color.go

@@ -152,17 +152,5 @@ func colorize(s string, escapeSeq string) string {
 	if !colorsEnabled || escapeSeq == "" {
 		return s
 	}
-
-	out := ""
-	if !strings.HasPrefix(s, EscapeStart) {
-		out += escapeSeq
-	}
-	out += strings.Replace(s, EscapeReset, EscapeReset+escapeSeq, -1)
-	if !strings.HasSuffix(out, EscapeReset) {
-		out += EscapeReset
-	}
-	if strings.Contains(out, escapeSeq+EscapeReset) {
-		out = strings.Replace(out, escapeSeq+EscapeReset, "", -1)
-	}
-	return out
+	return Escape(s, escapeSeq)
 }