Add support for colored legend (#47)

* Add support of legends for colored graphs --------- Co-authored-by: Rohit Gupta <rohitgtech+git@gmail.com>
guptarohit · Mar 24, 2024 · 27408ca · 27408ca
1 parent e398e0f
commit 27408ca
Show file tree

Hide file tree

Showing 6 changed files with 187 additions and 13 deletions.
diff --git a/README.md b/README.md
@@ -7,15 +7,15 @@ Go package to make lightweight ASCII line graphs ╭┈╯.
 ![image][]
 
 ## Installation
-``` bash
+```bash
 go get -u github.com/guptarohit/asciigraph@latest
 ```
 
 ## Usage
 
 ### Basic graph
 
-``` go
+```go
 package main
 
 import (
@@ -32,7 +32,7 @@ func main() {
 ```
 
 Running this example would render the following graph:
-``` bash
+```bash
   10.00 ┤        ╭╮
    9.00 ┤ ╭╮     ││
    8.00 ┤ ││   ╭╮││
@@ -46,7 +46,7 @@ Running this example would render the following graph:
 
 ### Multiple Series
 
-``` go
+```go
 package main
 
 import (
@@ -63,7 +63,7 @@ func main() {
 ```
 
 Running this example would render the following graph:
-``` bash
+```bash
  6.00 ┤    ╭─
  5.00 ┼╮   │
  4.00 ┤╰╮ ╭╯
@@ -75,7 +75,7 @@ Running this example would render the following graph:
 
 ### Colored graphs
 
-``` go
+```go
 package main
 
 import (
@@ -110,18 +110,54 @@ Running this example would render the following graph:
 
 ![colored_graph_image][]
 
+### Legends for colored graphs
+
+The graph can include legends for each series, making it easier to interpret.
+
+```go
+package main
+
+import (
+	"fmt"
+	"github.com/guptarohit/asciigraph"
+	"math"
+)
+
+func main() {
+	data := make([][]float64, 3)
+	for i := 0; i < 3; i++ {
+		for x := -12; x <= 12; x++ {
+			v := math.NaN()
+			if r := 12 - 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.Green, asciigraph.Blue),
+		asciigraph.SeriesLegends("Red", "Green", "Blue"),
+		asciigraph.Caption("Series with legends"))
+	fmt.Println(graph)
+}
+```
+Running this example would render the following graph:
+
+![graph_with_legends_image][]
+
 
 ## CLI Installation
 
 This package also brings a small utility for command line usage.
 
 Assuming `$GOPATH/bin` is in your `$PATH`, install CLI with following command:
-``` bash
+```bash
 go install github.com/guptarohit/asciigraph/cmd/asciigraph@latest
 ```
 
 or pull Docker image:
-``` bash
+```bash
 docker pull ghcr.io/guptarohit/asciigraph:latest
 ```
 
@@ -130,7 +166,7 @@ or download binaries from the [releases][] page.
 
 ## CLI Usage
 
-``` bash                                                                                                                ✘ 0|125  16:19:23
+```bash                                                                                                                ✘ 0|125  16:19:23
 > asciigraph --help
 Usage of asciigraph:
   asciigraph [options]
@@ -168,18 +204,18 @@ asciigraph expects data points from stdin. Invalid values are logged to stderr.
 
 
 Feed it data points via stdin:
-``` bash
+```bash
 seq 1 72 | asciigraph -h 10 -c "plot data from stdin"
 ```
 
 or use Docker image:
-``` bash
+```bash
 seq 1 72 | docker run -i --rm ghcr.io/guptarohit/asciigraph -h 10 -c "plot data from stdin"
 ```
 
 Output:
 
-``` bash
+```bash
  72.00 ┤                                                                  ╭────
  64.90 ┤                                                           ╭──────╯
  57.80 ┤                                                    ╭──────╯
@@ -195,7 +231,7 @@ Output:
 ```
 
 Example of **real-time graph** for data points stream via stdin:
-``` bash
+```bash
 ping -i.2 google.com | grep -oP '(?<=time=).*(?=ms)' --line-buffered | asciigraph -r -h 10 -w 40 -c "realtime plot data (google ping in ms) from stdin"
 ```
 [![asciinema][]][7]
@@ -229,3 +265,4 @@ Feel free to make a pull request! :octocat:
 [asciichart]: https://github.com/kroitor/asciichart
 [asciinema]: https://asciinema.org/a/382383.svg
 [7]: https://asciinema.org/a/382383
+[graph_with_legends_image]: https://github.com/guptarohit/asciigraph/assets/7895001/4066ee95-55ca-42a4-8a03-e73ce20df5d3
diff --git a/asciigraph.go b/asciigraph.go
@@ -253,5 +253,9 @@ func PlotMany(data [][]float64, options ...Option) string {
 		}
 	}
 
+	if len(config.SeriesLegends) > 0 {
+		addLegends(&lines, config, lenMax, config.Offset+maxWidth)
+	}
+
 	return lines.String()
 }
diff --git a/asciigraph_test.go b/asciigraph_test.go
@@ -358,6 +358,19 @@ func TestPlotMany(t *testing.T) {
  2.00 ┤\x1b[91m╭╭\x1b[0m
  1.00 ┤\x1b[91m││\x1b[0m
  0.00 ┼\x1b[91m╯╯\x1b[0m`},
+		{
+			[][]float64{{0, 1, 0}, {2, 3, 4, 3, 2}},
+			[]Option{SeriesColors(Red, Blue), SeriesLegends("Red", "Blue"),
+				Caption("legends with caption test")},
+			`
+ 4.00 ┤ [94m╭╮[0m
+ 3.00 ┤[94m╭╯╰╮[0m
+ 2.00 ┼[94m╯[0m  [94m╰[0m
+ 1.00 ┤[91m╭╮[0m
+ 0.00 ┼[91m╯╰[0m
+       legends with caption test
+
+       [91m■[0m Red   [94m■[0m Blue`},
 	}
 
 	for i := range cases {

diff --git a/examples/rainbow-legends/main.go b/examples/rainbow-legends/main.go
@@ -0,0 +1,67 @@
+package main
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/guptarohit/asciigraph"
+)
+
+func main() {
+	data := make([][]float64, 6)
+
+	// concentric semi-circles
+	for i := 0; i < 6; i++ {
+		for x := -40; x <= 40; x++ {
+			v := math.NaN()
+			if r := 40 - 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.Orange,
+		asciigraph.Yellow,
+		asciigraph.Green,
+		asciigraph.Blue,
+		asciigraph.Purple,
+	), asciigraph.SeriesLegends(
+		"Red",
+		"Orange",
+		"Yellow",
+		"Green",
+		"Blue",
+		"Purple",
+	),
+		asciigraph.Caption("Rainbow with color legends"))
+
+	fmt.Println(graph)
+	// Output:
+	//   20 ┤
+	//   20 ┤                               ╭───────╭╮───────╮
+	//   19 ┤                        ╭──╭───╭───────╭╮───────╮───╮──╮
+	//   18 ┤                    ╭─╭──╭─╭───╭───────╭╮───────╮───╮─╮──╮─╮
+	//   17 ┤                 ╭─╭─╭─╭─╭──╭──────────╯╰──────────╮──╮─╮─╮─╮─╮
+	//   16 ┤              ╭─╭─╭╭─╭─╭────╯                      ╰────╮─╮─╮╮─╮─╮
+	//   15 ┤            ╭╭─╭─╭╭─╭──╯                                ╰──╮─╮╮─╮─╮╮
+	//   14 ┤          ╭╭─╭╭─╭╭──╯                                      ╰──╮╮─╮╮─╮╮
+	//   13 ┤        ╭─╭╭╭─╭╭─╯                                            ╰─╮╮─╮╮╮─╮
+	//   12 ┤       ╭╭╭─╭╭╭─╯                                                ╰─╮╮╮─╮╮╮
+	//   11 ┤     ╭─╭╭╭╭╭─╯                                                    ╰─╮╮╮╮╮─╮
+	//   10 ┤    ╭╭─╭╭╭╭╯                                                        ╰╮╮╮╮─╮╮
+	//    9 ┤   ╭╭╯╭╭╭╭╯                                                          ╰╮╮╮╮╰╮╮
+	//    8 ┤  ╭╭╯╭╭╭╭╯                                                            ╰╮╮╮╮╰╮╮
+	//    7 ┤  ││╭╭╭╭╯                                                              ╰╮╮╮╮││
+	//    6 ┤ ╭╭╭╭╭╭╯                                                                ╰╮╮╮╮╮╮
+	//    5 ┤ ││││││                                                                  ││││││
+	//    4 ┤╭╭╭╭╭╭╯                                                                  ╰╮╮╮╮╮╮
+	//    3 ┤││││││                                                                    ││││││
+	//    2 ┤││││││                                                                    ││││││
+	//    1 ┤││││││                                                                    ││││││
+	//    0 ┼╶╶╶╶╶╯                                                                    ╰╴╴╴╴╴
+	//                                  Rainbow with color legends
+	//
+	//                   ■ Red   ■ Orange   ■ Yellow   ■ Green   ■ Blue   ■ Purple
+}
diff --git a/legend.go b/legend.go
@@ -0,0 +1,45 @@
+package asciigraph
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"unicode/utf8"
+)
+
+// Create legend item as a colored box and text
+func createLegendItem(text string, color AnsiColor) (string, int) {
+	return fmt.Sprintf(
+			"%s■%s %s",
+			color.String(),
+			Default.String(),
+			text,
+		),
+		// Can't use len() because of AnsiColor, add 2 for box and space
+		utf8.RuneCountInString(text) + 2
+}
+
+// Add legend for each series added to the graph
+func addLegends(lines *bytes.Buffer, config *config, lenMax int, leftPad int) {
+	lines.WriteString("\n\n")
+	lines.WriteString(strings.Repeat(" ", leftPad))
+
+	var legendsText string
+	var legendsTextLen int
+	rightPad := 3
+	for i, text := range config.SeriesLegends {
+		item, itemLen := createLegendItem(text, config.SeriesColors[i])
+		legendsText += item
+		legendsTextLen += itemLen
+
+		if i < len(config.SeriesLegends)-1 {
+			legendsText += strings.Repeat(" ", rightPad)
+			legendsTextLen += rightPad
+		}
+	}
+
+	if legendsTextLen < lenMax {
+		lines.WriteString(strings.Repeat(" ", (lenMax-legendsTextLen)/2))
+	}
+	lines.WriteString(legendsText)
+}
diff --git a/options.go b/options.go
@@ -20,6 +20,7 @@ type config struct {
 	AxisColor              AnsiColor
 	LabelColor             AnsiColor
 	SeriesColors           []AnsiColor
+	SeriesLegends          []string
 }
 
 // An optionFunc applies an option.
@@ -116,3 +117,10 @@ func SeriesColors(ac ...AnsiColor) Option {
 		c.SeriesColors = ac
 	})
 }
+
+// SeriesLegends sets the legend text for the corresponding series.
+func SeriesLegends(text ...string) Option {
+	return optionFunc(func(c *config) {
+		c.SeriesLegends = text
+	})
+}