1package color
2
3import (
4	"fmt"
5	"io"
6	"os"
7	"strconv"
8	"strings"
9	"sync"
10
11	"github.com/mattn/go-colorable"
12	"github.com/mattn/go-isatty"
13)
14
15var (
16	// NoColor defines if the output is colorized or not. It's dynamically set to
17	// false or true based on the stdout's file descriptor referring to a terminal
18	// or not. It's also set to true if the NO_COLOR environment variable is
19	// set (regardless of its value). This is a global option and affects all
20	// colors. For more control over each color block use the methods
21	// DisableColor() individually.
22	NoColor = noColorExists() || os.Getenv("TERM") == "dumb" ||
23		(!isatty.IsTerminal(os.Stdout.Fd()) && !isatty.IsCygwinTerminal(os.Stdout.Fd()))
24
25	// Output defines the standard output of the print functions. By default
26	// os.Stdout is used.
27	Output = colorable.NewColorableStdout()
28
29	// Error defines a color supporting writer for os.Stderr.
30	Error = colorable.NewColorableStderr()
31
32	// colorsCache is used to reduce the count of created Color objects and
33	// allows to reuse already created objects with required Attribute.
34	colorsCache   = make(map[Attribute]*Color)
35	colorsCacheMu sync.Mutex // protects colorsCache
36)
37
38// noColorExists returns true if the environment variable NO_COLOR exists.
39func noColorExists() bool {
40	_, exists := os.LookupEnv("NO_COLOR")
41	return exists
42}
43
44// Color defines a custom color object which is defined by SGR parameters.
45type Color struct {
46	params  []Attribute
47	noColor *bool
48}
49
50// Attribute defines a single SGR Code
51type Attribute int
52
53const escape = "\x1b"
54
55// Base attributes
56const (
57	Reset Attribute = iota
58	Bold
59	Faint
60	Italic
61	Underline
62	BlinkSlow
63	BlinkRapid
64	ReverseVideo
65	Concealed
66	CrossedOut
67)
68
69// Foreground text colors
70const (
71	FgBlack Attribute = iota + 30
72	FgRed
73	FgGreen
74	FgYellow
75	FgBlue
76	FgMagenta
77	FgCyan
78	FgWhite
79)
80
81// Foreground Hi-Intensity text colors
82const (
83	FgHiBlack Attribute = iota + 90
84	FgHiRed
85	FgHiGreen
86	FgHiYellow
87	FgHiBlue
88	FgHiMagenta
89	FgHiCyan
90	FgHiWhite
91)
92
93// Background text colors
94const (
95	BgBlack Attribute = iota + 40
96	BgRed
97	BgGreen
98	BgYellow
99	BgBlue
100	BgMagenta
101	BgCyan
102	BgWhite
103)
104
105// Background Hi-Intensity text colors
106const (
107	BgHiBlack Attribute = iota + 100
108	BgHiRed
109	BgHiGreen
110	BgHiYellow
111	BgHiBlue
112	BgHiMagenta
113	BgHiCyan
114	BgHiWhite
115)
116
117// New returns a newly created color object.
118func New(value ...Attribute) *Color {
119	c := &Color{
120		params: make([]Attribute, 0),
121	}
122
123	if noColorExists() {
124		c.noColor = boolPtr(true)
125	}
126
127	c.Add(value...)
128	return c
129}
130
131// Set sets the given parameters immediately. It will change the color of
132// output with the given SGR parameters until color.Unset() is called.
133func Set(p ...Attribute) *Color {
134	c := New(p...)
135	c.Set()
136	return c
137}
138
139// Unset resets all escape attributes and clears the output. Usually should
140// be called after Set().
141func Unset() {
142	if NoColor {
143		return
144	}
145
146	fmt.Fprintf(Output, "%s[%dm", escape, Reset)
147}
148
149// Set sets the SGR sequence.
150func (c *Color) Set() *Color {
151	if c.isNoColorSet() {
152		return c
153	}
154
155	fmt.Fprintf(Output, c.format())
156	return c
157}
158
159func (c *Color) unset() {
160	if c.isNoColorSet() {
161		return
162	}
163
164	Unset()
165}
166
167func (c *Color) setWriter(w io.Writer) *Color {
168	if c.isNoColorSet() {
169		return c
170	}
171
172	fmt.Fprintf(w, c.format())
173	return c
174}
175
176func (c *Color) unsetWriter(w io.Writer) {
177	if c.isNoColorSet() {
178		return
179	}
180
181	if NoColor {
182		return
183	}
184
185	fmt.Fprintf(w, "%s[%dm", escape, Reset)
186}
187
188// Add is used to chain SGR parameters. Use as many as parameters to combine
189// and create custom color objects. Example: Add(color.FgRed, color.Underline).
190func (c *Color) Add(value ...Attribute) *Color {
191	c.params = append(c.params, value...)
192	return c
193}
194
195func (c *Color) prepend(value Attribute) {
196	c.params = append(c.params, 0)
197	copy(c.params[1:], c.params[0:])
198	c.params[0] = value
199}
200
201// Fprint formats using the default formats for its operands and writes to w.
202// Spaces are added between operands when neither is a string.
203// It returns the number of bytes written and any write error encountered.
204// On Windows, users should wrap w with colorable.NewColorable() if w is of
205// type *os.File.
206func (c *Color) Fprint(w io.Writer, a ...interface{}) (n int, err error) {
207	c.setWriter(w)
208	defer c.unsetWriter(w)
209
210	return fmt.Fprint(w, a...)
211}
212
213// Print formats using the default formats for its operands and writes to
214// standard output. Spaces are added between operands when neither is a
215// string. It returns the number of bytes written and any write error
216// encountered. This is the standard fmt.Print() method wrapped with the given
217// color.
218func (c *Color) Print(a ...interface{}) (n int, err error) {
219	c.Set()
220	defer c.unset()
221
222	return fmt.Fprint(Output, a...)
223}
224
225// Fprintf formats according to a format specifier and writes to w.
226// It returns the number of bytes written and any write error encountered.
227// On Windows, users should wrap w with colorable.NewColorable() if w is of
228// type *os.File.
229func (c *Color) Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
230	c.setWriter(w)
231	defer c.unsetWriter(w)
232
233	return fmt.Fprintf(w, format, a...)
234}
235
236// Printf formats according to a format specifier and writes to standard output.
237// It returns the number of bytes written and any write error encountered.
238// This is the standard fmt.Printf() method wrapped with the given color.
239func (c *Color) Printf(format string, a ...interface{}) (n int, err error) {
240	c.Set()
241	defer c.unset()
242
243	return fmt.Fprintf(Output, format, a...)
244}
245
246// Fprintln formats using the default formats for its operands and writes to w.
247// Spaces are always added between operands and a newline is appended.
248// On Windows, users should wrap w with colorable.NewColorable() if w is of
249// type *os.File.
250func (c *Color) Fprintln(w io.Writer, a ...interface{}) (n int, err error) {
251	c.setWriter(w)
252	defer c.unsetWriter(w)
253
254	return fmt.Fprintln(w, a...)
255}
256
257// Println formats using the default formats for its operands and writes to
258// standard output. Spaces are always added between operands and a newline is
259// appended. It returns the number of bytes written and any write error
260// encountered. This is the standard fmt.Print() method wrapped with the given
261// color.
262func (c *Color) Println(a ...interface{}) (n int, err error) {
263	c.Set()
264	defer c.unset()
265
266	return fmt.Fprintln(Output, a...)
267}
268
269// Sprint is just like Print, but returns a string instead of printing it.
270func (c *Color) Sprint(a ...interface{}) string {
271	return c.wrap(fmt.Sprint(a...))
272}
273
274// Sprintln is just like Println, but returns a string instead of printing it.
275func (c *Color) Sprintln(a ...interface{}) string {
276	return c.wrap(fmt.Sprintln(a...))
277}
278
279// Sprintf is just like Printf, but returns a string instead of printing it.
280func (c *Color) Sprintf(format string, a ...interface{}) string {
281	return c.wrap(fmt.Sprintf(format, a...))
282}
283
284// FprintFunc returns a new function that prints the passed arguments as
285// colorized with color.Fprint().
286func (c *Color) FprintFunc() func(w io.Writer, a ...interface{}) {
287	return func(w io.Writer, a ...interface{}) {
288		c.Fprint(w, a...)
289	}
290}
291
292// PrintFunc returns a new function that prints the passed arguments as
293// colorized with color.Print().
294func (c *Color) PrintFunc() func(a ...interface{}) {
295	return func(a ...interface{}) {
296		c.Print(a...)
297	}
298}
299
300// FprintfFunc returns a new function that prints the passed arguments as
301// colorized with color.Fprintf().
302func (c *Color) FprintfFunc() func(w io.Writer, format string, a ...interface{}) {
303	return func(w io.Writer, format string, a ...interface{}) {
304		c.Fprintf(w, format, a...)
305	}
306}
307
308// PrintfFunc returns a new function that prints the passed arguments as
309// colorized with color.Printf().
310func (c *Color) PrintfFunc() func(format string, a ...interface{}) {
311	return func(format string, a ...interface{}) {
312		c.Printf(format, a...)
313	}
314}
315
316// FprintlnFunc returns a new function that prints the passed arguments as
317// colorized with color.Fprintln().
318func (c *Color) FprintlnFunc() func(w io.Writer, a ...interface{}) {
319	return func(w io.Writer, a ...interface{}) {
320		c.Fprintln(w, a...)
321	}
322}
323
324// PrintlnFunc returns a new function that prints the passed arguments as
325// colorized with color.Println().
326func (c *Color) PrintlnFunc() func(a ...interface{}) {
327	return func(a ...interface{}) {
328		c.Println(a...)
329	}
330}
331
332// SprintFunc returns a new function that returns colorized strings for the
333// given arguments with fmt.Sprint(). Useful to put into or mix into other
334// string. Windows users should use this in conjunction with color.Output, example:
335//
336//	put := New(FgYellow).SprintFunc()
337//	fmt.Fprintf(color.Output, "This is a %s", put("warning"))
338func (c *Color) SprintFunc() func(a ...interface{}) string {
339	return func(a ...interface{}) string {
340		return c.wrap(fmt.Sprint(a...))
341	}
342}
343
344// SprintfFunc returns a new function that returns colorized strings for the
345// given arguments with fmt.Sprintf(). Useful to put into or mix into other
346// string. Windows users should use this in conjunction with color.Output.
347func (c *Color) SprintfFunc() func(format string, a ...interface{}) string {
348	return func(format string, a ...interface{}) string {
349		return c.wrap(fmt.Sprintf(format, a...))
350	}
351}
352
353// SprintlnFunc returns a new function that returns colorized strings for the
354// given arguments with fmt.Sprintln(). Useful to put into or mix into other
355// string. Windows users should use this in conjunction with color.Output.
356func (c *Color) SprintlnFunc() func(a ...interface{}) string {
357	return func(a ...interface{}) string {
358		return c.wrap(fmt.Sprintln(a...))
359	}
360}
361
362// sequence returns a formatted SGR sequence to be plugged into a "\x1b[...m"
363// an example output might be: "1;36" -> bold cyan
364func (c *Color) sequence() string {
365	format := make([]string, len(c.params))
366	for i, v := range c.params {
367		format[i] = strconv.Itoa(int(v))
368	}
369
370	return strings.Join(format, ";")
371}
372
373// wrap wraps the s string with the colors attributes. The string is ready to
374// be printed.
375func (c *Color) wrap(s string) string {
376	if c.isNoColorSet() {
377		return s
378	}
379
380	return c.format() + s + c.unformat()
381}
382
383func (c *Color) format() string {
384	return fmt.Sprintf("%s[%sm", escape, c.sequence())
385}
386
387func (c *Color) unformat() string {
388	return fmt.Sprintf("%s[%dm", escape, Reset)
389}
390
391// DisableColor disables the color output. Useful to not change any existing
392// code and still being able to output. Can be used for flags like
393// "--no-color". To enable back use EnableColor() method.
394func (c *Color) DisableColor() {
395	c.noColor = boolPtr(true)
396}
397
398// EnableColor enables the color output. Use it in conjunction with
399// DisableColor(). Otherwise this method has no side effects.
400func (c *Color) EnableColor() {
401	c.noColor = boolPtr(false)
402}
403
404func (c *Color) isNoColorSet() bool {
405	// check first if we have user set action
406	if c.noColor != nil {
407		return *c.noColor
408	}
409
410	// if not return the global option, which is disabled by default
411	return NoColor
412}
413
414// Equals returns a boolean value indicating whether two colors are equal.
415func (c *Color) Equals(c2 *Color) bool {
416	if len(c.params) != len(c2.params) {
417		return false
418	}
419
420	for _, attr := range c.params {
421		if !c2.attrExists(attr) {
422			return false
423		}
424	}
425
426	return true
427}
428
429func (c *Color) attrExists(a Attribute) bool {
430	for _, attr := range c.params {
431		if attr == a {
432			return true
433		}
434	}
435
436	return false
437}
438
439func boolPtr(v bool) *bool {
440	return &v
441}
442
443func getCachedColor(p Attribute) *Color {
444	colorsCacheMu.Lock()
445	defer colorsCacheMu.Unlock()
446
447	c, ok := colorsCache[p]
448	if !ok {
449		c = New(p)
450		colorsCache[p] = c
451	}
452
453	return c
454}
455
456func colorPrint(format string, p Attribute, a ...interface{}) {
457	c := getCachedColor(p)
458
459	if !strings.HasSuffix(format, "\n") {
460		format += "\n"
461	}
462
463	if len(a) == 0 {
464		c.Print(format)
465	} else {
466		c.Printf(format, a...)
467	}
468}
469
470func colorString(format string, p Attribute, a ...interface{}) string {
471	c := getCachedColor(p)
472
473	if len(a) == 0 {
474		return c.SprintFunc()(format)
475	}
476
477	return c.SprintfFunc()(format, a...)
478}
479
480// Black is a convenient helper function to print with black foreground. A
481// newline is appended to format by default.
482func Black(format string, a ...interface{}) { colorPrint(format, FgBlack, a...) }
483
484// Red is a convenient helper function to print with red foreground. A
485// newline is appended to format by default.
486func Red(format string, a ...interface{}) { colorPrint(format, FgRed, a...) }
487
488// Green is a convenient helper function to print with green foreground. A
489// newline is appended to format by default.
490func Green(format string, a ...interface{}) { colorPrint(format, FgGreen, a...) }
491
492// Yellow is a convenient helper function to print with yellow foreground.
493// A newline is appended to format by default.
494func Yellow(format string, a ...interface{}) { colorPrint(format, FgYellow, a...) }
495
496// Blue is a convenient helper function to print with blue foreground. A
497// newline is appended to format by default.
498func Blue(format string, a ...interface{}) { colorPrint(format, FgBlue, a...) }
499
500// Magenta is a convenient helper function to print with magenta foreground.
501// A newline is appended to format by default.
502func Magenta(format string, a ...interface{}) { colorPrint(format, FgMagenta, a...) }
503
504// Cyan is a convenient helper function to print with cyan foreground. A
505// newline is appended to format by default.
506func Cyan(format string, a ...interface{}) { colorPrint(format, FgCyan, a...) }
507
508// White is a convenient helper function to print with white foreground. A
509// newline is appended to format by default.
510func White(format string, a ...interface{}) { colorPrint(format, FgWhite, a...) }
511
512// BlackString is a convenient helper function to return a string with black
513// foreground.
514func BlackString(format string, a ...interface{}) string { return colorString(format, FgBlack, a...) }
515
516// RedString is a convenient helper function to return a string with red
517// foreground.
518func RedString(format string, a ...interface{}) string { return colorString(format, FgRed, a...) }
519
520// GreenString is a convenient helper function to return a string with green
521// foreground.
522func GreenString(format string, a ...interface{}) string { return colorString(format, FgGreen, a...) }
523
524// YellowString is a convenient helper function to return a string with yellow
525// foreground.
526func YellowString(format string, a ...interface{}) string { return colorString(format, FgYellow, a...) }
527
528// BlueString is a convenient helper function to return a string with blue
529// foreground.
530func BlueString(format string, a ...interface{}) string { return colorString(format, FgBlue, a...) }
531
532// MagentaString is a convenient helper function to return a string with magenta
533// foreground.
534func MagentaString(format string, a ...interface{}) string {
535	return colorString(format, FgMagenta, a...)
536}
537
538// CyanString is a convenient helper function to return a string with cyan
539// foreground.
540func CyanString(format string, a ...interface{}) string { return colorString(format, FgCyan, a...) }
541
542// WhiteString is a convenient helper function to return a string with white
543// foreground.
544func WhiteString(format string, a ...interface{}) string { return colorString(format, FgWhite, a...) }
545
546// HiBlack is a convenient helper function to print with hi-intensity black foreground. A
547// newline is appended to format by default.
548func HiBlack(format string, a ...interface{}) { colorPrint(format, FgHiBlack, a...) }
549
550// HiRed is a convenient helper function to print with hi-intensity red foreground. A
551// newline is appended to format by default.
552func HiRed(format string, a ...interface{}) { colorPrint(format, FgHiRed, a...) }
553
554// HiGreen is a convenient helper function to print with hi-intensity green foreground. A
555// newline is appended to format by default.
556func HiGreen(format string, a ...interface{}) { colorPrint(format, FgHiGreen, a...) }
557
558// HiYellow is a convenient helper function to print with hi-intensity yellow foreground.
559// A newline is appended to format by default.
560func HiYellow(format string, a ...interface{}) { colorPrint(format, FgHiYellow, a...) }
561
562// HiBlue is a convenient helper function to print with hi-intensity blue foreground. A
563// newline is appended to format by default.
564func HiBlue(format string, a ...interface{}) { colorPrint(format, FgHiBlue, a...) }
565
566// HiMagenta is a convenient helper function to print with hi-intensity magenta foreground.
567// A newline is appended to format by default.
568func HiMagenta(format string, a ...interface{}) { colorPrint(format, FgHiMagenta, a...) }
569
570// HiCyan is a convenient helper function to print with hi-intensity cyan foreground. A
571// newline is appended to format by default.
572func HiCyan(format string, a ...interface{}) { colorPrint(format, FgHiCyan, a...) }
573
574// HiWhite is a convenient helper function to print with hi-intensity white foreground. A
575// newline is appended to format by default.
576func HiWhite(format string, a ...interface{}) { colorPrint(format, FgHiWhite, a...) }
577
578// HiBlackString is a convenient helper function to return a string with hi-intensity black
579// foreground.
580func HiBlackString(format string, a ...interface{}) string {
581	return colorString(format, FgHiBlack, a...)
582}
583
584// HiRedString is a convenient helper function to return a string with hi-intensity red
585// foreground.
586func HiRedString(format string, a ...interface{}) string { return colorString(format, FgHiRed, a...) }
587
588// HiGreenString is a convenient helper function to return a string with hi-intensity green
589// foreground.
590func HiGreenString(format string, a ...interface{}) string {
591	return colorString(format, FgHiGreen, a...)
592}
593
594// HiYellowString is a convenient helper function to return a string with hi-intensity yellow
595// foreground.
596func HiYellowString(format string, a ...interface{}) string {
597	return colorString(format, FgHiYellow, a...)
598}
599
600// HiBlueString is a convenient helper function to return a string with hi-intensity blue
601// foreground.
602func HiBlueString(format string, a ...interface{}) string { return colorString(format, FgHiBlue, a...) }
603
604// HiMagentaString is a convenient helper function to return a string with hi-intensity magenta
605// foreground.
606func HiMagentaString(format string, a ...interface{}) string {
607	return colorString(format, FgHiMagenta, a...)
608}
609
610// HiCyanString is a convenient helper function to return a string with hi-intensity cyan
611// foreground.
612func HiCyanString(format string, a ...interface{}) string { return colorString(format, FgHiCyan, a...) }
613
614// HiWhiteString is a convenient helper function to return a string with hi-intensity white
615// foreground.
616func HiWhiteString(format string, a ...interface{}) string {
617	return colorString(format, FgHiWhite, a...)
618}
619