color.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511
  1. package color
  2. import (
  3. "fmt"
  4. "io"
  5. "os"
  6. "strconv"
  7. "strings"
  8. "sync"
  9. "github.com/mattn/go-colorable"
  10. "github.com/mattn/go-isatty"
  11. )
  12. var (
  13. // NoColor defines if the output is colorized or not. It's dynamically set to
  14. // false or true based on the stdout's file descriptor referring to a terminal
  15. // or not. This is a global option and affects all colors. For more control
  16. // over each color block use the methods DisableColor() individually.
  17. NoColor = !isatty.IsTerminal(os.Stdout.Fd()) || os.Getenv("TERM") == "dumb"
  18. // Output defines the standard output of the print functions. By default
  19. // os.Stdout is used.
  20. Output = colorable.NewColorableStdout()
  21. // colorsCache is used to reduce the count of created Color objects and
  22. // allows to reuse already created objects with required Attribute.
  23. colorsCache = make(map[Attribute]*Color)
  24. colorsCacheMu sync.Mutex // protects colorsCache
  25. )
  26. // Color defines a custom color object which is defined by SGR parameters.
  27. type Color struct {
  28. params []Attribute
  29. noColor *bool
  30. }
  31. // Attribute defines a single SGR Code
  32. type Attribute int
  33. const escape = "\x1b"
  34. // Base attributes
  35. const (
  36. Reset Attribute = iota
  37. Bold
  38. Faint
  39. Italic
  40. Underline
  41. BlinkSlow
  42. BlinkRapid
  43. ReverseVideo
  44. Concealed
  45. CrossedOut
  46. )
  47. // Foreground text colors
  48. const (
  49. FgBlack Attribute = iota + 30
  50. FgRed
  51. FgGreen
  52. FgYellow
  53. FgBlue
  54. FgMagenta
  55. FgCyan
  56. FgWhite
  57. )
  58. // Foreground Hi-Intensity text colors
  59. const (
  60. FgHiBlack Attribute = iota + 90
  61. FgHiRed
  62. FgHiGreen
  63. FgHiYellow
  64. FgHiBlue
  65. FgHiMagenta
  66. FgHiCyan
  67. FgHiWhite
  68. )
  69. // Background text colors
  70. const (
  71. BgBlack Attribute = iota + 40
  72. BgRed
  73. BgGreen
  74. BgYellow
  75. BgBlue
  76. BgMagenta
  77. BgCyan
  78. BgWhite
  79. )
  80. // Background Hi-Intensity text colors
  81. const (
  82. BgHiBlack Attribute = iota + 100
  83. BgHiRed
  84. BgHiGreen
  85. BgHiYellow
  86. BgHiBlue
  87. BgHiMagenta
  88. BgHiCyan
  89. BgHiWhite
  90. )
  91. // New returns a newly created color object.
  92. func New(value ...Attribute) *Color {
  93. c := &Color{params: make([]Attribute, 0)}
  94. c.Add(value...)
  95. return c
  96. }
  97. // Set sets the given parameters immediately. It will change the color of
  98. // output with the given SGR parameters until color.Unset() is called.
  99. func Set(p ...Attribute) *Color {
  100. c := New(p...)
  101. c.Set()
  102. return c
  103. }
  104. // Unset resets all escape attributes and clears the output. Usually should
  105. // be called after Set().
  106. func Unset() {
  107. if NoColor {
  108. return
  109. }
  110. fmt.Fprintf(Output, "%s[%dm", escape, Reset)
  111. }
  112. // Set sets the SGR sequence.
  113. func (c *Color) Set() *Color {
  114. if c.isNoColorSet() {
  115. return c
  116. }
  117. fmt.Fprintf(Output, c.format())
  118. return c
  119. }
  120. func (c *Color) unset() {
  121. if c.isNoColorSet() {
  122. return
  123. }
  124. Unset()
  125. }
  126. func (c *Color) setWriter(w io.Writer) *Color {
  127. if c.isNoColorSet() {
  128. return c
  129. }
  130. fmt.Fprintf(w, c.format())
  131. return c
  132. }
  133. func (c *Color) unsetWriter(w io.Writer) {
  134. if c.isNoColorSet() {
  135. return
  136. }
  137. if NoColor {
  138. return
  139. }
  140. fmt.Fprintf(w, "%s[%dm", escape, Reset)
  141. }
  142. // Add is used to chain SGR parameters. Use as many as parameters to combine
  143. // and create custom color objects. Example: Add(color.FgRed, color.Underline).
  144. func (c *Color) Add(value ...Attribute) *Color {
  145. c.params = append(c.params, value...)
  146. return c
  147. }
  148. func (c *Color) prepend(value Attribute) {
  149. c.params = append(c.params, 0)
  150. copy(c.params[1:], c.params[0:])
  151. c.params[0] = value
  152. }
  153. // Fprint formats using the default formats for its operands and writes to w.
  154. // Spaces are added between operands when neither is a string.
  155. // It returns the number of bytes written and any write error encountered.
  156. // On Windows, users should wrap w with colorable.NewColorable() if w is of
  157. // type *os.File.
  158. func (c *Color) Fprint(w io.Writer, a ...interface{}) (n int, err error) {
  159. c.setWriter(w)
  160. defer c.unsetWriter(w)
  161. return fmt.Fprint(w, a...)
  162. }
  163. // Print formats using the default formats for its operands and writes to
  164. // standard output. Spaces are added between operands when neither is a
  165. // string. It returns the number of bytes written and any write error
  166. // encountered. This is the standard fmt.Print() method wrapped with the given
  167. // color.
  168. func (c *Color) Print(a ...interface{}) (n int, err error) {
  169. c.Set()
  170. defer c.unset()
  171. return fmt.Fprint(Output, a...)
  172. }
  173. // Fprintf formats according to a format specifier and writes to w.
  174. // It returns the number of bytes written and any write error encountered.
  175. // On Windows, users should wrap w with colorable.NewColorable() if w is of
  176. // type *os.File.
  177. func (c *Color) Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
  178. c.setWriter(w)
  179. defer c.unsetWriter(w)
  180. return fmt.Fprintf(w, format, a...)
  181. }
  182. // Printf formats according to a format specifier and writes to standard output.
  183. // It returns the number of bytes written and any write error encountered.
  184. // This is the standard fmt.Printf() method wrapped with the given color.
  185. func (c *Color) Printf(format string, a ...interface{}) (n int, err error) {
  186. c.Set()
  187. defer c.unset()
  188. return fmt.Fprintf(Output, format, a...)
  189. }
  190. // Fprintln formats using the default formats for its operands and writes to w.
  191. // Spaces are always added between operands and a newline is appended.
  192. // On Windows, users should wrap w with colorable.NewColorable() if w is of
  193. // type *os.File.
  194. func (c *Color) Fprintln(w io.Writer, a ...interface{}) (n int, err error) {
  195. c.setWriter(w)
  196. defer c.unsetWriter(w)
  197. return fmt.Fprintln(w, a...)
  198. }
  199. // Println formats using the default formats for its operands and writes to
  200. // standard output. Spaces are always added between operands and a newline is
  201. // appended. It returns the number of bytes written and any write error
  202. // encountered. This is the standard fmt.Print() method wrapped with the given
  203. // color.
  204. func (c *Color) Println(a ...interface{}) (n int, err error) {
  205. c.Set()
  206. defer c.unset()
  207. return fmt.Fprintln(Output, a...)
  208. }
  209. // FprintFunc returns a new function that prints the passed arguments as
  210. // colorized with color.Fprint().
  211. func (c *Color) FprintFunc() func(w io.Writer, a ...interface{}) {
  212. return func(w io.Writer, a ...interface{}) {
  213. c.Fprint(w, a...)
  214. }
  215. }
  216. // PrintFunc returns a new function that prints the passed arguments as
  217. // colorized with color.Print().
  218. func (c *Color) PrintFunc() func(a ...interface{}) {
  219. return func(a ...interface{}) {
  220. c.Print(a...)
  221. }
  222. }
  223. // FprintfFunc returns a new function that prints the passed arguments as
  224. // colorized with color.Fprintf().
  225. func (c *Color) FprintfFunc() func(w io.Writer, format string, a ...interface{}) {
  226. return func(w io.Writer, format string, a ...interface{}) {
  227. c.Fprintf(w, format, a...)
  228. }
  229. }
  230. // PrintfFunc returns a new function that prints the passed arguments as
  231. // colorized with color.Printf().
  232. func (c *Color) PrintfFunc() func(format string, a ...interface{}) {
  233. return func(format string, a ...interface{}) {
  234. c.Printf(format, a...)
  235. }
  236. }
  237. // FprintlnFunc returns a new function that prints the passed arguments as
  238. // colorized with color.Fprintln().
  239. func (c *Color) FprintlnFunc() func(w io.Writer, a ...interface{}) {
  240. return func(w io.Writer, a ...interface{}) {
  241. c.Fprintln(w, a...)
  242. }
  243. }
  244. // PrintlnFunc returns a new function that prints the passed arguments as
  245. // colorized with color.Println().
  246. func (c *Color) PrintlnFunc() func(a ...interface{}) {
  247. return func(a ...interface{}) {
  248. c.Println(a...)
  249. }
  250. }
  251. // SprintFunc returns a new function that returns colorized strings for the
  252. // given arguments with fmt.Sprint(). Useful to put into or mix into other
  253. // string. Windows users should use this in conjunction with color.Output, example:
  254. //
  255. // put := New(FgYellow).SprintFunc()
  256. // fmt.Fprintf(color.Output, "This is a %s", put("warning"))
  257. func (c *Color) SprintFunc() func(a ...interface{}) string {
  258. return func(a ...interface{}) string {
  259. return c.wrap(fmt.Sprint(a...))
  260. }
  261. }
  262. // SprintfFunc returns a new function that returns colorized strings for the
  263. // given arguments with fmt.Sprintf(). Useful to put into or mix into other
  264. // string. Windows users should use this in conjunction with color.Output.
  265. func (c *Color) SprintfFunc() func(format string, a ...interface{}) string {
  266. return func(format string, a ...interface{}) string {
  267. return c.wrap(fmt.Sprintf(format, a...))
  268. }
  269. }
  270. // SprintlnFunc returns a new function that returns colorized strings for the
  271. // given arguments with fmt.Sprintln(). Useful to put into or mix into other
  272. // string. Windows users should use this in conjunction with color.Output.
  273. func (c *Color) SprintlnFunc() func(a ...interface{}) string {
  274. return func(a ...interface{}) string {
  275. return c.wrap(fmt.Sprintln(a...))
  276. }
  277. }
  278. // sequence returns a formated SGR sequence to be plugged into a "\x1b[...m"
  279. // an example output might be: "1;36" -> bold cyan
  280. func (c *Color) sequence() string {
  281. format := make([]string, len(c.params))
  282. for i, v := range c.params {
  283. format[i] = strconv.Itoa(int(v))
  284. }
  285. return strings.Join(format, ";")
  286. }
  287. // wrap wraps the s string with the colors attributes. The string is ready to
  288. // be printed.
  289. func (c *Color) wrap(s string) string {
  290. if c.isNoColorSet() {
  291. return s
  292. }
  293. return c.format() + s + c.unformat()
  294. }
  295. func (c *Color) format() string {
  296. return fmt.Sprintf("%s[%sm", escape, c.sequence())
  297. }
  298. func (c *Color) unformat() string {
  299. return fmt.Sprintf("%s[%dm", escape, Reset)
  300. }
  301. // DisableColor disables the color output. Useful to not change any existing
  302. // code and still being able to output. Can be used for flags like
  303. // "--no-color". To enable back use EnableColor() method.
  304. func (c *Color) DisableColor() {
  305. c.noColor = boolPtr(true)
  306. }
  307. // EnableColor enables the color output. Use it in conjunction with
  308. // DisableColor(). Otherwise this method has no side effects.
  309. func (c *Color) EnableColor() {
  310. c.noColor = boolPtr(false)
  311. }
  312. func (c *Color) isNoColorSet() bool {
  313. // check first if we have user setted action
  314. if c.noColor != nil {
  315. return *c.noColor
  316. }
  317. // if not return the global option, which is disabled by default
  318. return NoColor
  319. }
  320. // Equals returns a boolean value indicating whether two colors are equal.
  321. func (c *Color) Equals(c2 *Color) bool {
  322. if len(c.params) != len(c2.params) {
  323. return false
  324. }
  325. for _, attr := range c.params {
  326. if !c2.attrExists(attr) {
  327. return false
  328. }
  329. }
  330. return true
  331. }
  332. func (c *Color) attrExists(a Attribute) bool {
  333. for _, attr := range c.params {
  334. if attr == a {
  335. return true
  336. }
  337. }
  338. return false
  339. }
  340. func boolPtr(v bool) *bool {
  341. return &v
  342. }
  343. func getCachedColor(p Attribute) *Color {
  344. colorsCacheMu.Lock()
  345. defer colorsCacheMu.Unlock()
  346. c, ok := colorsCache[p]
  347. if !ok {
  348. c = New(p)
  349. colorsCache[p] = c
  350. }
  351. return c
  352. }
  353. func colorPrint(format string, p Attribute, a ...interface{}) {
  354. c := getCachedColor(p)
  355. if !strings.HasSuffix(format, "\n") {
  356. format += "\n"
  357. }
  358. if len(a) == 0 {
  359. c.Print(format)
  360. } else {
  361. c.Printf(format, a...)
  362. }
  363. }
  364. func colorString(format string, p Attribute, a ...interface{}) string {
  365. c := getCachedColor(p)
  366. if len(a) == 0 {
  367. return c.SprintFunc()(format)
  368. }
  369. return c.SprintfFunc()(format, a...)
  370. }
  371. // Black is an convenient helper function to print with black foreground. A
  372. // newline is appended to format by default.
  373. func Black(format string, a ...interface{}) { colorPrint(format, FgBlack, a...) }
  374. // Red is an convenient helper function to print with red foreground. A
  375. // newline is appended to format by default.
  376. func Red(format string, a ...interface{}) { colorPrint(format, FgRed, a...) }
  377. // Green is an convenient helper function to print with green foreground. A
  378. // newline is appended to format by default.
  379. func Green(format string, a ...interface{}) { colorPrint(format, FgGreen, a...) }
  380. // Yellow is an convenient helper function to print with yellow foreground.
  381. // A newline is appended to format by default.
  382. func Yellow(format string, a ...interface{}) { colorPrint(format, FgYellow, a...) }
  383. // Blue is an convenient helper function to print with blue foreground. A
  384. // newline is appended to format by default.
  385. func Blue(format string, a ...interface{}) { colorPrint(format, FgBlue, a...) }
  386. // Magenta is an convenient helper function to print with magenta foreground.
  387. // A newline is appended to format by default.
  388. func Magenta(format string, a ...interface{}) { colorPrint(format, FgMagenta, a...) }
  389. // Cyan is an convenient helper function to print with cyan foreground. A
  390. // newline is appended to format by default.
  391. func Cyan(format string, a ...interface{}) { colorPrint(format, FgCyan, a...) }
  392. // White is an convenient helper function to print with white foreground. A
  393. // newline is appended to format by default.
  394. func White(format string, a ...interface{}) { colorPrint(format, FgWhite, a...) }
  395. // BlackString is an convenient helper function to return a string with black
  396. // foreground.
  397. func BlackString(format string, a ...interface{}) string { return colorString(format, FgBlack, a...) }
  398. // RedString is an convenient helper function to return a string with red
  399. // foreground.
  400. func RedString(format string, a ...interface{}) string { return colorString(format, FgRed, a...) }
  401. // GreenString is an convenient helper function to return a string with green
  402. // foreground.
  403. func GreenString(format string, a ...interface{}) string { return colorString(format, FgGreen, a...) }
  404. // YellowString is an convenient helper function to return a string with yellow
  405. // foreground.
  406. func YellowString(format string, a ...interface{}) string { return colorString(format, FgYellow, a...) }
  407. // BlueString is an convenient helper function to return a string with blue
  408. // foreground.
  409. func BlueString(format string, a ...interface{}) string { return colorString(format, FgBlue, a...) }
  410. // MagentaString is an convenient helper function to return a string with magenta
  411. // foreground.
  412. func MagentaString(format string, a ...interface{}) string {
  413. return colorString(format, FgMagenta, a...)
  414. }
  415. // CyanString is an convenient helper function to return a string with cyan
  416. // foreground.
  417. func CyanString(format string, a ...interface{}) string { return colorString(format, FgCyan, a...) }
  418. // WhiteString is an convenient helper function to return a string with white
  419. // foreground.
  420. func WhiteString(format string, a ...interface{}) string { return colorString(format, FgWhite, a...) }