You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

527 lines
14 KiB

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