From ad2adb2933210a19b8ec9884105f6cac8bc97aa7 Mon Sep 17 00:00:00 2001
From: Jason Walton <jwalton@solinkcorp.com>
Date: Tue, 22 Mar 2022 13:33:01 -0400
Subject: [PATCH] feat: Add ColorFn convenience type.

fix #3
---
 colorModels.go | 32 ++++++++++++++++----------------
 gchalk.go      | 11 +++++++----
 gchalk_test.go | 15 +++++++++++----
 3 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/colorModels.go b/colorModels.go
index 7cb81fa..cced518 100644
--- a/colorModels.go
+++ b/colorModels.go
@@ -5,7 +5,7 @@ import (
 )
 
 // Ansi returns a function which colors a string using the ANSI 16 color pallette.
-func Ansi(code uint8) func(strs ...string) string {
+func Ansi(code uint8) ColorFn {
 	return rootBuilder.WithAnsi(code).applyStyle
 }
 
@@ -15,7 +15,7 @@ func WithAnsi(code uint8) *Builder {
 }
 
 // Ansi returns a function which colors a string using the ANSI 16 color pallette.
-func (builder *Builder) Ansi(code uint8) func(strs ...string) string {
+func (builder *Builder) Ansi(code uint8) ColorFn {
 	return builder.WithAnsi(code).applyStyle
 }
 
@@ -29,7 +29,7 @@ func (builder *Builder) WithAnsi(code uint8) *Builder {
 }
 
 // BgAnsi returns a function which colors the background of a string using the ANSI 16 color pallette.
-func BgAnsi(code uint8) func(strs ...string) string {
+func BgAnsi(code uint8) ColorFn {
 	return rootBuilder.WithBgAnsi(code).applyStyle
 }
 
@@ -39,7 +39,7 @@ func WithBgAnsi(code uint8) *Builder {
 }
 
 // BgAnsi returns a function which colors the background of a string using the ANSI 16 color pallette.
-func (builder *Builder) BgAnsi(code uint8) func(strs ...string) string {
+func (builder *Builder) BgAnsi(code uint8) ColorFn {
 	return builder.WithBgAnsi(code).applyStyle
 }
 
@@ -51,7 +51,7 @@ func (builder *Builder) WithBgAnsi(code uint8) *Builder {
 // Ansi256 returns a function which colors a string using the ANSI 256 color pallette.
 // If ANSI 256 color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func Ansi256(code uint8) func(strs ...string) string {
+func Ansi256(code uint8) ColorFn {
 	return rootBuilder.WithAnsi256(code).applyStyle
 }
 
@@ -65,7 +65,7 @@ func WithAnsi256(code uint8) *Builder {
 // Ansi256 returns a function which colors a string using the ANSI 256 color pallette.
 // If ANSI 256 color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func (builder *Builder) Ansi256(code uint8) func(strs ...string) string {
+func (builder *Builder) Ansi256(code uint8) ColorFn {
 	return builder.WithAnsi256(code).applyStyle
 }
 
@@ -83,7 +83,7 @@ func (builder *Builder) WithAnsi256(code uint8) *Builder {
 // BgAnsi256 returns a function which colors the background of a string using the ANSI 256 color pallette.
 // If ANSI 256 color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func BgAnsi256(code uint8) func(strs ...string) string {
+func BgAnsi256(code uint8) ColorFn {
 	return rootBuilder.WithBgAnsi256(code).applyStyle
 }
 
@@ -97,7 +97,7 @@ func WithBgAnsi256(code uint8) *Builder {
 // BgAnsi256 returns a function which colors the background of a string using the ANSI 256 color pallette.
 // If ANSI 256 color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func (builder *Builder) BgAnsi256(code uint8) func(strs ...string) string {
+func (builder *Builder) BgAnsi256(code uint8) ColorFn {
 	return builder.WithBgAnsi256(code).applyStyle
 }
 
@@ -115,7 +115,7 @@ func (builder *Builder) WithBgAnsi256(code uint8) *Builder {
 // RGB returns a function which colors a string using true color support.  Note that
 // if true color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func RGB(r uint8, g uint8, b uint8) func(strs ...string) string {
+func RGB(r uint8, g uint8, b uint8) ColorFn {
 	return rootBuilder.WithRGB(r, g, b).applyStyle
 }
 
@@ -129,7 +129,7 @@ func WithRGB(r uint8, g uint8, b uint8) *Builder {
 // RGB returns a function which colors a string using true color support.  Note that
 // if true color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func (builder *Builder) RGB(r uint8, g uint8, b uint8) func(strs ...string) string {
+func (builder *Builder) RGB(r uint8, g uint8, b uint8) ColorFn {
 	return builder.WithRGB(r, g, b).applyStyle
 }
 
@@ -146,7 +146,7 @@ func (builder *Builder) WithRGB(r uint8, g uint8, b uint8) *Builder {
 // BgRGB returns a function which colors the background of a string using true color support.  Note that
 // if true color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func BgRGB(r uint8, g uint8, b uint8) func(strs ...string) string {
+func BgRGB(r uint8, g uint8, b uint8) ColorFn {
 	return rootBuilder.WithBgRGB(r, g, b).applyStyle
 }
 
@@ -160,7 +160,7 @@ func WithBgRGB(r uint8, g uint8, b uint8) *Builder {
 // BgRGB returns a function which colors the background of a string using true color support.  Note that
 // if true color support is unavailable, this will automatically convert the color
 // to the closest available color.
-func (builder *Builder) BgRGB(r uint8, g uint8, b uint8) func(strs ...string) string {
+func (builder *Builder) BgRGB(r uint8, g uint8, b uint8) ColorFn {
 	return builder.WithBgRGB(r, g, b).applyStyle
 }
 
@@ -178,7 +178,7 @@ func (builder *Builder) WithBgRGB(r uint8, g uint8, b uint8) *Builder {
 // hexadecimal color string (e.g. "#FF00FF" or "#FFF"). If true color support is
 // unavailable, this will automatically convert the color  to the closest
 // available color.
-func Hex(hex string) func(strs ...string) string {
+func Hex(hex string) ColorFn {
 	return rootBuilder.WithRGB(ansistyles.HexToRGB(hex)).applyStyle
 }
 
@@ -193,7 +193,7 @@ func WithHex(hex string) *Builder {
 // hexadecimal color string (e.g. "#FF00FF" or "#FFF"). If true color support is
 // unavailable, this will automatically convert the color  to the closest
 // available color.
-func (builder *Builder) Hex(hex string) func(strs ...string) string {
+func (builder *Builder) Hex(hex string) ColorFn {
 	return builder.WithRGB(ansistyles.HexToRGB(hex)).applyStyle
 }
 
@@ -208,7 +208,7 @@ func (builder *Builder) WithHex(hex string) *Builder {
 // hexadecimal color string (e.g. "#FF00FF" or "#FFF"). If true color support is
 // unavailable, this will automatically convert the color  to the closest
 // available color.
-func BgHex(hex string) func(strs ...string) string {
+func BgHex(hex string) ColorFn {
 	return rootBuilder.WithBgRGB(ansistyles.HexToRGB(hex)).applyStyle
 }
 
@@ -223,7 +223,7 @@ func WithBgHex(hex string) *Builder {
 // hexadecimal color string (e.g. "#FF00FF" or "#FFF"). If true color support is
 // unavailable, this will automatically convert the color  to the closest
 // available color.
-func (builder *Builder) BgHex(hex string) func(strs ...string) string {
+func (builder *Builder) BgHex(hex string) ColorFn {
 	return builder.WithBgRGB(ansistyles.HexToRGB(hex)).applyStyle
 }
 
diff --git a/gchalk.go b/gchalk.go
index 838f47d..790553f 100644
--- a/gchalk.go
+++ b/gchalk.go
@@ -52,6 +52,9 @@ import (
 // ColorLevel represents the ANSI color level supported by the terminal.
 type ColorLevel = supportscolor.ColorLevel
 
+// ColorFn is a convenience type for a function that takes in a string and returns a colored string.
+type ColorFn func(str ...string) string
+
 const (
 	// LevelNone represents a terminal that does not support color at all.
 	LevelNone ColorLevel = supportscolor.None
@@ -283,7 +286,7 @@ func (builder *Builder) GetLevel() ColorLevel {
 // styles. Styles can be specified as a named style (e.g. "red", "bgRed", "bgred"),
 // or as a hex color ("#ff00ff" or "bg#ff00ff").  If the style cannot
 // be parsed, this will panic.
-func StyleMust(styles ...string) func(strs ...string) string {
+func StyleMust(styles ...string) ColorFn {
 	return rootBuilder.WithStyleMust(styles...).applyStyle
 }
 
@@ -299,7 +302,7 @@ func WithStyleMust(styles ...string) *Builder {
 // styles. Styles can be specified as a named style (e.g. "red", "bgRed", "bgred"),
 // or as a hex color ("#ff00ff" or "bg#ff00ff").  If the style cannot
 // be parsed, this will panic.
-func (builder *Builder) StyleMust(styles ...string) func(strs ...string) string {
+func (builder *Builder) StyleMust(styles ...string) ColorFn {
 	return builder.WithStyleMust(styles...).applyStyle
 }
 
@@ -320,7 +323,7 @@ func (builder *Builder) WithStyleMust(styles ...string) *Builder {
 // styles. Styles can be specified as a named style (e.g. "red", "bgRed", "bgred"),
 // or as a hex color ("#ff00ff" or "bg#ff00ff").  If the style cannot
 // be parsed, this will return an error.
-func Style(styles ...string) (func(strs ...string) string, error) {
+func Style(styles ...string) (ColorFn, error) {
 	newBuilder, err := rootBuilder.WithStyle(styles...)
 	if err != nil {
 		return rootBuilder.applyStyle, err
@@ -340,7 +343,7 @@ func WithStyle(styles ...string) (*Builder, error) {
 // styles. Styles can be specified as a named style (e.g. "red", "bgRed", "bgred"),
 // or as a hex color ("#ff00ff" or "bg#ff00ff").  If the style cannot
 // be parsed, this will return an error.
-func (builder *Builder) Style(styles ...string) (func(strs ...string) string, error) {
+func (builder *Builder) Style(styles ...string) (ColorFn, error) {
 	newBuilder, err := rootBuilder.WithStyle(styles...)
 	if err != nil {
 		return rootBuilder.applyStyle, err
diff --git a/gchalk_test.go b/gchalk_test.go
index 1608ebe..bd42a50 100644
--- a/gchalk_test.go
+++ b/gchalk_test.go
@@ -68,10 +68,11 @@ func TestResetAllStyles(t *testing.T) {
 func TestCachingMultipleStyles(t *testing.T) {
 	gchalk := New(ForceLevel(LevelAnsi16m))
 
-	red := gchalk.WithRed().Red
-	green := gchalk.WithRed().Green
-	redBold := gchalk.WithRed().WithRed().Bold
-	greenBold := gchalk.WithRed().WithGreen().Bold
+	var red, green, redBold, greenBold ColorFn
+	red = gchalk.WithRed().Red
+	green = gchalk.WithRed().Green
+	redBold = gchalk.WithRed().WithRed().Bold
+	greenBold = gchalk.WithRed().WithGreen().Bold
 
 	if red("foo") == green("foo") {
 		t.Errorf("red and green should produce different output")
@@ -223,6 +224,12 @@ func TestStyle(t *testing.T) {
 	str := styler("foo")
 	assertEqual(t, str, "foo")
 	assertEqual(t, fmt.Sprintf("%v", err), "No such style: idonotexist")
+
+	var styler2 ColorFn = gchalk.StyleMust("red")
+	styler2("a")
+	// Make sure we can assign a ColorFn to a regular func.
+	var styler3 func(s ...string) string = gchalk.StyleMust("blue")
+	styler3("b")
 }
 
 func TestStyleWithHex(t *testing.T) {