Merge pull request #226 from polymorcodeus/mart1082/optsetcustomcolor

Added OptionSetCustomColorCodes to pass in custom color codes.

Author: schollz

progressbar.go

@@ -79,6 +79,8 @@ type config struct {
 
 	// whether the output is expected to contain color codes
 	colorCodes bool
+	// custom colors to use for colorCodes
+	customColors map[string]string
 
 	// show rate of change in kB/sec or MB/sec
 	showBytes bool
@@ -280,6 +282,14 @@ func OptionEnableColorCodes(colorCodes bool) Option {
 	}
 }
 
+// OptionSetCutomColorCodes overrides DefaultColors for color codes
+// using mitchellh/colorstring.Colorize in func customColorstringColor
+func OptionSetCustomColorCodes(customColors map[string]string) Option {
+	return func(p *ProgressBar) {
+		p.config.customColors = customColors
+	}
+}
+
 // OptionSetElapsedTime will enable elapsed time. Always enabled if OptionSetPredictTime is true.
 func OptionSetElapsedTime(elapsedTime bool) Option {
 	return func(p *ProgressBar) {
@@ -1069,13 +1079,24 @@ func (p *ProgressBar) StartHTTPServer(hostPort string) *http.Server {
 	return server
 }
 
+// use customstring.Colorize to customize color palette
+func customColorstringColor(col map[string]string, str string) string {
+	cs := colorstring.Colorize{Colors: col, Reset: true}
+	return cs.Color(str)
+}
+
 // regex matching ansi escape codes
 var ansiRegex = regexp.MustCompile(`\x1b\[[0-9;]*[a-zA-Z]`)
 
 func getStringWidth(c config, str string) int {
 	if c.colorCodes {
 		// convert any color codes in the progress bar into the respective ANSI codes
-		str = colorstring.Color(str)
+		// if customColors have been option set - create cs with Colorize
+		if len(c.customColors) > 0 {
+			str = customColorstringColor(c.customColors, str)
+		} else {
+			str = colorstring.Color(str)
+		}
 	}
 
 	// the width of the string, if printed to the console
@@ -1362,7 +1383,12 @@ func renderProgressBar(c config, s *state) (int, error) {
 
 	if c.colorCodes {
 		// convert any color codes in the progress bar into the respective ANSI codes
-		str = colorstring.Color(str)
+		// if customColors have been option set - create cs with Colorize
+		if len(c.customColors) > 0 {
+			str = customColorstringColor(c.customColors, str)
+		} else {
+			str = colorstring.Color(str)
+		}
 	}
 
 	if c.useANSICodes {