Merge pull request #35 from NathanBaulch/ansi-colors

Ansi colors

README.md

@@ -73,6 +73,43 @@ Running this example would render the following graph:
  0.00 ┼╯     ╰
 ```
 
+### Colored graphs
+
+``` go
+package main
+
+import (
+    "fmt"
+    "github.com/guptarohit/asciigraph"
+)
+
+func main() {
+	data := make([][]float64, 4)
+
+	for i := 0; i < 4; i++ {
+		for x := -20; x <= 20; x++ {
+			v := math.NaN()
+			if r := 20 - i; x >= -r && x <= r {
+				v = math.Sqrt(math.Pow(float64(r), 2)-math.Pow(float64(x), 2)) / 2
+			}
+			data[i] = append(data[i], v)
+		}
+	}
+	graph := asciigraph.PlotMany(data, asciigraph.Precision(0), asciigraph.SeriesColors(
+		asciigraph.Red,
+		asciigraph.Yellow,
+		asciigraph.Green,
+		asciigraph.Blue,
+	))
+
+	fmt.Println(graph)
+}
+```
+
+Running this example would render the following graph:
+
+![colored_graph_image][]
+
 ## Command line interface
 
 This package also brings a small utility for command line usage.
@@ -81,20 +118,28 @@ This package also brings a small utility for command line usage.
 Usage of asciigraph:
   asciigraph [options]
 Options:
+  -ac axis color
+    	y-axis color of the plot
   -b buffer
     	data points buffer when realtime graph enabled, default equal to `width`
   -c caption
     	caption for the graph
+  -cc caption color
+    	caption color of the plot
   -f fps
     	set fps to control how frequently graph to be rendered when realtime graph enabled (default 24)
   -h height
     	height in text rows, 0 for auto-scaling
+  -lc label color
+    	y-axis label color of the plot
   -o offset
     	offset in columns, for the label (default 3)
   -p precision
     	precision of data point labels along the y-axis (default 2)
   -r realtime
     	enables realtime graph for data stream
+  -sc series color
+    	series color of the plot
   -w width
     	width in columns, 0 for auto-scaling
 asciigraph expects data points from stdin. Invalid values are logged to stderr.
@@ -174,6 +219,7 @@ Feel free to make a pull request! :octocat:
 [Mentioned in Awesome Go]: https://awesome.re/mentioned-badge-flat.svg
 [6]: https://github.com/avelino/awesome-go#advanced-console-uis
 [image]: https://user-images.githubusercontent.com/7895001/41509956-b1b2b3d0-7279-11e8-9d19-d7dea17d5e44.png
+[colored_graph_image]: https://user-images.githubusercontent.com/7895001/166443444-40ad8113-2c0f-46d7-9c75-1cf08435ce15.png
 [releases]: https://github.com/guptarohit/asciigraph/releases
 [asciichart]: https://github.com/kroitor/asciichart
 [asciinema]: https://asciinema.org/a/382383.svg

asciigraph.go

@@ -77,13 +77,18 @@ func PlotMany(data [][]float64, options ...Option) string {
 	rows := int(math.Abs(float64(intmax2 - intmin2)))
 	width := lenMax + config.Offset
 
-	plot := make([][]string, rows+1)
+	type cell struct {
+		Text  string
+		Color AnsiColor
+	}
+	plot := make([][]cell, rows+1)
 
 	// initialise empty 2D grid
 	for i := 0; i < rows+1; i++ {
-		line := make([]string, width)
+		line := make([]cell, width)
 		for j := 0; j < width; j++ {
-			line[j] = " "
+			line[j].Text = " "
+			line[j].Color = Default
 		}
 		plot[i] = line
 	}
@@ -123,17 +128,26 @@ func PlotMany(data [][]float64, options ...Option) string {
 		w := y - intmin2
 		h := int(math.Max(float64(config.Offset)-float64(len(label)), 0))
 
-		plot[w][h] = label
-		plot[w][config.Offset-1] = "┤"
+		plot[w][h].Text = label
+		plot[w][h].Color = config.LabelColor
+		plot[w][config.Offset-1].Text = "┤"
+		plot[w][config.Offset-1].Color = config.AxisColor
 	}
 
 	for i := range data {
 		series := data[i]
+
+		color := Default
+		if i < len(config.SeriesColors) {
+			color = config.SeriesColors[i]
+		}
+
 		var y0, y1 int
 
 		if !math.IsNaN(series[0]) {
 			y0 = int(round(series[0]*ratio) - min2)
-			plot[rows-y0][config.Offset-1] = "┼" // first value
+			plot[rows-y0][config.Offset-1].Text = "┼" // first value
+			plot[rows-y0][config.Offset-1].Color = config.AxisColor
 		}
 
 		for x := 0; x < len(series)-1; x++ { // plot the line
@@ -146,36 +160,44 @@ func PlotMany(data [][]float64, options ...Option) string {
 
 			if math.IsNaN(d1) && !math.IsNaN(d0) {
 				y0 = int(round(d0*ratio) - float64(intmin2))
-				plot[rows-y0][x+config.Offset] = "╴"
+				plot[rows-y0][x+config.Offset].Text = "╴"
+				plot[rows-y0][x+config.Offset].Color = color
 				continue
 			}
 
 			if math.IsNaN(d0) && !math.IsNaN(d1) {
 				y1 = int(round(d1*ratio) - float64(intmin2))
-				plot[rows-y1][x+config.Offset] = "╶"
+				plot[rows-y1][x+config.Offset].Text = "╶"
+				plot[rows-y1][x+config.Offset].Color = color
 				continue
 			}
 
 			y0 = int(round(d0*ratio) - float64(intmin2))
 			y1 = int(round(d1*ratio) - float64(intmin2))
 
 			if y0 == y1 {
-				plot[rows-y0][x+config.Offset] = "─"
+				plot[rows-y0][x+config.Offset].Text = "─"
 			} else {
 				if y0 > y1 {
-					plot[rows-y1][x+config.Offset] = "╰"
-					plot[rows-y0][x+config.Offset] = "╮"
+					plot[rows-y1][x+config.Offset].Text = "╰"
+					plot[rows-y0][x+config.Offset].Text = "╮"
 				} else {
-					plot[rows-y1][x+config.Offset] = "╭"
-					plot[rows-y0][x+config.Offset] = "╯"
+					plot[rows-y1][x+config.Offset].Text = "╭"
+					plot[rows-y0][x+config.Offset].Text = "╯"
 				}
 
 				start := int(math.Min(float64(y0), float64(y1))) + 1
 				end := int(math.Max(float64(y0), float64(y1)))
 				for y := start; y < end; y++ {
-					plot[rows-y][x+config.Offset] = "│"
+					plot[rows-y][x+config.Offset].Text = "│"
 				}
 			}
+
+			start := int(math.Min(float64(y0), float64(y1)))
+			end := int(math.Max(float64(y0), float64(y1)))
+			for y := start; y <= end; y++ {
+				plot[rows-y][x+config.Offset].Color = color
+			}
 		}
 	}
 
@@ -189,14 +211,23 @@ func PlotMany(data [][]float64, options ...Option) string {
 		// remove trailing spaces
 		lastCharIndex := 0
 		for i := width - 1; i >= 0; i-- {
-			if horizontal[i] != " " {
+			if horizontal[i].Text != " " {
 				lastCharIndex = i
 				break
 			}
 		}
 
+		c := Default
 		for _, v := range horizontal[:lastCharIndex+1] {
-			lines.WriteString(v)
+			if v.Color != c {
+				c = v.Color
+				lines.WriteString(c.String())
+			}
+
+			lines.WriteString(v.Text)
+		}
+		if c != Default {
+			lines.WriteString(Default.String())
 		}
 	}
 
@@ -207,7 +238,13 @@ func PlotMany(data [][]float64, options ...Option) string {
 		if len(config.Caption) < lenMax {
 			lines.WriteString(strings.Repeat(" ", (lenMax-len(config.Caption))/2))
 		}
+		if config.CaptionColor != Default {
+			lines.WriteString(config.CaptionColor.String())
+		}
 		lines.WriteString(config.Caption)
+		if config.CaptionColor != Default {
+			lines.WriteString(Default.String())
+		}
 	}
 
 	return lines.String()

asciigraph_test.go

@@ -242,13 +242,20 @@ func TestPlot(t *testing.T) {
  -0.025449 ┤                         ╰─╮         ╭─╯
  -0.033940 ┤                           ╰╮   ╭────╯
  -0.042430 ┤                            ╰───╯`},
+
+		{
+			[]float64{math.NaN(), 1},
+			[]Option{Caption("color test"), CaptionColor(Red), AxisColor(Green), LabelColor(Blue)},
+			`
+\x1b[94m 1.00\x1b[0m \x1b[32m┤\x1b[0m╶
+       \x1b[91mcolor test\x1b[0m`},
 	}
 
 	for i := range cases {
 		name := fmt.Sprintf("%d", i)
 		t.Run(name, func(t *testing.T) {
 			c := cases[i]
-			expected := strings.TrimPrefix(c.expected, "\n")
+			expected := strings.Replace(strings.TrimPrefix(c.expected, "\n"), `\x1b`, "\x1b", -1)
 			actual := Plot(c.data, c.opts...)
 			if actual != expected {
 				conf := configure(config{}, c.opts)
@@ -297,13 +304,29 @@ func TestPlotMany(t *testing.T) {
  1.00 ┤ ╭───╮
  0.00 ┼─╯   ╰╴
         interpolation test`},
+
+		{
+			[][]float64{{0, 0}, {math.NaN(), 0}},
+			[]Option{SeriesColors(Red)},
+			" 0.00 ┼╶"},
+		{
+			[][]float64{{0, 0}, {math.NaN(), 0}},
+			[]Option{SeriesColors(Default, Red)},
+			" 0.00 ┼\x1b[91m╶\x1b[0m"},
+		{
+			[][]float64{{math.NaN(), 0, 2}, {0, 2}},
+			[]Option{SeriesColors(Red, Red)},
+			`
+ 2.00 ┤\x1b[91m╭╭\x1b[0m
+ 1.00 ┤\x1b[91m││\x1b[0m
+ 0.00 ┼\x1b[91m╯╯\x1b[0m`},
 	}
 
 	for i := range cases {
 		name := fmt.Sprintf("%d", i)
 		t.Run(name, func(t *testing.T) {
 			c := cases[i]
-			expected := strings.TrimPrefix(c.expected, "\n")
+			expected := strings.Replace(strings.TrimPrefix(c.expected, "\n"), `\x1b`, "\x1b", -1)
 			actual := PlotMany(c.data, c.opts...)
 			if actual != expected {
 				conf := configure(config{}, c.opts)

cmd/asciigraph/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"errors"
 	"flag"
 	"fmt"
 	"log"
@@ -21,6 +22,10 @@ var (
 	enableRealTime     bool
 	realTimeDataBuffer int
 	fps                float64 = 24
+	seriesColor        asciigraph.AnsiColor
+	captionColor       asciigraph.AnsiColor
+	axisColor          asciigraph.AnsiColor
+	labelColor         asciigraph.AnsiColor
 )
 
 func main() {
@@ -39,6 +44,42 @@ func main() {
 	flag.BoolVar(&enableRealTime, "r", enableRealTime, "enables `realtime` graph for data stream")
 	flag.IntVar(&realTimeDataBuffer, "b", realTimeDataBuffer, "data points `buffer` when realtime graph enabled, default equal to `width`")
 	flag.Float64Var(&fps, "f", fps, "set `fps` to control how frequently graph to be rendered when realtime graph enabled")
+	flag.Func("sc", "`series color` of the plot", func(str string) error {
+		if c, ok := asciigraph.ColorNames[str]; !ok {
+			return errors.New("unrecognized color, check available color names at https://www.w3.org/TR/SVG11/types.html#ColorKeywords")
+		} else {
+			seriesColor = c
+			return nil
+		}
+	})
+
+	flag.Func("cc", "`caption color` of the plot", func(str string) error {
+		if c, ok := asciigraph.ColorNames[str]; !ok {
+			return errors.New("unrecognized color, check available color names at https://www.w3.org/TR/SVG11/types.html#ColorKeywords")
+		} else {
+			captionColor = c
+			return nil
+		}
+	})
+
+	flag.Func("ac", "y-`axis color` of the plot", func(str string) error {
+		if c, ok := asciigraph.ColorNames[str]; !ok {
+			return errors.New("unrecognized color, check available color names at https://www.w3.org/TR/SVG11/types.html#ColorKeywords")
+		} else {
+			axisColor = c
+			return nil
+		}
+	})
+
+	flag.Func("lc", "y-axis `label color` of the plot", func(str string) error {
+		if c, ok := asciigraph.ColorNames[str]; !ok {
+			return errors.New("unrecognized color, check available color names at https://www.w3.org/TR/SVG11/types.html#ColorKeywords")
+		} else {
+			labelColor = c
+			return nil
+		}
+	})
+
 	flag.Parse()
 
 	data := make([]float64, 0, 64)
@@ -73,7 +114,12 @@ func main() {
 					asciigraph.Width(int(width)),
 					asciigraph.Offset(int(offset)),
 					asciigraph.Precision(precision),
-					asciigraph.Caption(caption))
+					asciigraph.Caption(caption),
+					asciigraph.SeriesColors(seriesColor),
+					asciigraph.CaptionColor(captionColor),
+					asciigraph.AxisColor(axisColor),
+					asciigraph.LabelColor(labelColor),
+				)
 				asciigraph.Clear()
 				fmt.Println(plot)
 				nextFlushTime = time.Now().Add(flushInterval)
@@ -94,7 +140,12 @@ func main() {
 			asciigraph.Width(int(width)),
 			asciigraph.Offset(int(offset)),
 			asciigraph.Precision(precision),
-			asciigraph.Caption(caption))
+			asciigraph.Caption(caption),
+			asciigraph.SeriesColors(seriesColor),
+			asciigraph.CaptionColor(captionColor),
+			asciigraph.AxisColor(axisColor),
+			asciigraph.LabelColor(labelColor),
+		)
 
 		fmt.Println(plot)
 	}