[tools] Extract unicode progress bar to library

Change-Id: Ia0710777e7a133b3e335de917364858dd887618d
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/132640
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

tools/src/cmd/run-cts/main.go

@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
-	"math"
 	"net/http"
 	"os"
 	"os/exec"
@@ -42,6 +41,7 @@ import (
 	"dawn.googlesource.com/dawn/tools/src/cov"
 	"dawn.googlesource.com/dawn/tools/src/fileutils"
 	"dawn.googlesource.com/dawn/tools/src/git"
+	"dawn.googlesource.com/dawn/tools/src/progressbar"
 	"github.com/mattn/go-colorable"
 	"github.com/mattn/go-isatty"
 )
@@ -911,7 +911,7 @@ func (r *runner) streamResults(ctx context.Context, wg *sync.WaitGroup, results
 	// Helper function for printing a progress bar.
 	lastStatusUpdate, animFrame := time.Now(), 0
 	updateProgress := func() {
-		fmt.Fprint(r.stdout, ansiProgressBar(animFrame, numTests, numByExpectedStatus))
+		drawProgressBar(r.stdout, animFrame, numTests, numByExpectedStatus)
 		animFrame++
 		lastStatusUpdate = time.Now()
 	}
@@ -971,7 +971,7 @@ func (r *runner) streamResults(ctx context.Context, wg *sync.WaitGroup, results
 			covTree.Add(SplitCTSQuery(res.testcase), res.coverage)
 		}
 	}
-	fmt.Fprint(r.stdout, ansiProgressBar(animFrame, numTests, numByExpectedStatus))
+	drawProgressBar(r.stdout, animFrame, numTests, numByExpectedStatus)
 
 	// All done. Print final stats.
 	fmt.Fprintf(r.stdout, "\nCompleted in %v\n", timeTaken)
@@ -1101,6 +1101,14 @@ var statusColor = map[status]string{
 	fail:    red,
 }
 
+var pbStatusColor = map[status]progressbar.Color{
+	pass:    progressbar.Green,
+	warn:    progressbar.Yellow,
+	skip:    progressbar.Cyan,
+	timeout: progressbar.Yellow,
+	fail:    progressbar.Red,
+}
+
 // expectedStatus is a test status, along with a boolean to indicate whether the
 // status matches the test expectations
 type expectedStatus struct {
@@ -1274,69 +1282,26 @@ func alignRight(val interface{}, width int) string {
 	return strings.Repeat(" ", padding) + s
 }
 
-// ansiProgressBar returns a string with an ANSI-colored progress bar, providing
+// drawProgressBar draws an ANSI-colored progress bar, providing realtime
-// realtime information about the status of the CTS run.
+// information about the status of the CTS run.
 // Note: We'll want to skip this if !isatty or if we're running on windows.
-func ansiProgressBar(animFrame int, numTests int, numByExpectedStatus map[expectedStatus]int) string {
+func drawProgressBar(out io.Writer, animFrame int, numTests int, numByExpectedStatus map[expectedStatus]int) {
-	const barWidth = 50
+	bar := progressbar.Status{Total: numTests}
-
-	animSymbols := []rune{'⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'}
-	blockSymbols := []rune{'▏', '▎', '▍', '▌', '▋', '▊', '▉'}
-
-	numBlocksPrinted := 0
-
-	buf := &strings.Builder{}
-	fmt.Fprint(buf, string(animSymbols[animFrame%len(animSymbols)]), " [")
-	animFrame++
-
-	numFinished := 0
-
 	for _, status := range statuses {
 		for _, expected := range []bool{true, false} {
-			color := statusColor[status]
+			if num := numByExpectedStatus[expectedStatus{status, expected}]; num > 0 {
-			if expected {
+				bar.Segments = append(bar.Segments,
-				color += bold
+					progressbar.Segment{
-			}
+						Count:       num,
-
+						Color:       pbStatusColor[status],
-			num := numByExpectedStatus[expectedStatus{status, expected}]
+						Bold:        expected,
-			numFinished += num
+						Transparent: expected,
-			statusFrac := float64(num) / float64(numTests)
+					})
-			fNumBlocks := barWidth * statusFrac
-			fmt.Fprint(buf, color)
-			numBlocks := int(math.Ceil(fNumBlocks))
-			if expected {
-				if numBlocks > 1 {
-					fmt.Fprint(buf, strings.Repeat(string("░"), numBlocks))
-				}
-			} else {
-				if numBlocks > 1 {
-					fmt.Fprint(buf, strings.Repeat(string("▉"), numBlocks))
-				}
-				if numBlocks > 0 {
-					frac := fNumBlocks - math.Floor(fNumBlocks)
-					symbol := blockSymbols[int(math.Round(frac*float64(len(blockSymbols)-1)))]
-					fmt.Fprint(buf, string(symbol))
 			}
 		}
-			numBlocksPrinted += numBlocks
 	}
-	}
+	const width = 50
-
+	bar.Draw(out, width, colors, animFrame)
-	if barWidth > numBlocksPrinted {
-		fmt.Fprint(buf, strings.Repeat(string(" "), barWidth-numBlocksPrinted))
-	}
-	fmt.Fprint(buf, ansiReset)
-	fmt.Fprint(buf, "] ", percentage(numFinished, numTests))
-
-	if colors {
-		// move cursor to start of line so the bar is overridden
-		fmt.Fprint(buf, positionLeft)
-	} else {
-		// cannot move cursor, so newline
-		fmt.Fprintln(buf)
-	}
-
-	return buf.String()
 }
 
 // testcaseStatus is a pair of testcase name and result status

tools/src/progressbar/progressbar.go

@@ -0,0 +1,222 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package progressbar provides functions for drawing unicode progress bars to
+// the terminal
+package progressbar
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"math"
+	"strings"
+	"time"
+)
+
+// Defaults for the Config
+const (
+	DefaultRefreshRate = time.Millisecond * 100
+	DefaultWidth       = 50
+	DefaultANSIColors  = true
+)
+
+// Config holds configuration options for a ProgressBar
+type Config struct {
+	RefreshRate time.Duration
+	Width       int
+	ANSIColors  bool
+}
+
+// Color is an enumerator of colors
+type Color int
+
+// Color enumerators
+const (
+	White Color = iota
+	Red
+	Green
+	Yellow
+	Blue
+	Magenta
+	Cyan
+)
+
+// Segment describes a single segment of the ProgressBar
+type Segment struct {
+	Count       int
+	Color       Color
+	Transparent bool
+	Bold        bool
+}
+
+// Status holds the updated data of the ProgressBar
+type Status struct {
+	Total    int
+	Segments []Segment
+}
+
+// ProgressBar returns a string with an ANSI-colored progress bar, providing
+// realtime information about the status of the CTS run.
+// Note: We'll want to skip this if !isatty or if we're running on windows.
+type ProgressBar struct {
+	Config
+	out io.Writer
+	c   chan Status
+}
+
+// New returns a new ProgressBar that streams output to out.
+// Call ProgressBar.Stop() once finished.
+func New(out io.Writer, cfg *Config) *ProgressBar {
+	p := &ProgressBar{out: out, c: make(chan Status)}
+	if cfg != nil {
+		p.Config = *cfg
+	} else {
+		p.ANSIColors = DefaultANSIColors
+	}
+	if p.RefreshRate == 0 {
+		p.RefreshRate = DefaultRefreshRate
+	}
+	if p.Width == 0 {
+		p.Width = DefaultWidth
+	}
+	go func() {
+		var status *Status
+		t := time.NewTicker(p.RefreshRate)
+		defer t.Stop()
+		for frame := 0; ; frame++ {
+			select {
+			case s, ok := <-p.c:
+				if !ok {
+					return
+				}
+				status = &s
+			case <-t.C:
+				if status != nil {
+					status.Draw(out, p.Width, p.ANSIColors, frame)
+				}
+			}
+		}
+	}()
+	return p
+}
+
+// Update updates the ProgressBar with the given status
+func (p *ProgressBar) Update(s Status) {
+	p.c <- s
+}
+
+// Stop stops drawing the progress bar.
+// Once called, the ProgressBar must not be used.
+func (p *ProgressBar) Stop() {
+	close(p.c)
+}
+
+// Draw draws the ProgressBar status to out
+func (s Status) Draw(out io.Writer, width int, ansiColors bool, animFrame int) {
+	// ANSI escape sequences
+	const (
+		escape       = "\u001B["
+		positionLeft = escape + "0G"
+		ansiReset    = escape + "0m"
+
+		bold = escape + "1m"
+
+		red     = escape + "31m"
+		green   = escape + "32m"
+		yellow  = escape + "33m"
+		blue    = escape + "34m"
+		magenta = escape + "35m"
+		cyan    = escape + "36m"
+		white   = escape + "37m"
+	)
+
+	animSymbols := []rune{'⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'}
+	blockSymbols := []rune{'▏', '▎', '▍', '▌', '▋', '▊', '▉'}
+
+	numBlocksPrinted := 0
+
+	buf := &bytes.Buffer{}
+	fmt.Fprint(buf, "  ", string(animSymbols[animFrame%len(animSymbols)]), " [")
+
+	numFinished := 0
+	for _, seg := range s.Segments {
+		if ansiColors {
+			switch seg.Color {
+			case Red:
+				buf.WriteString(red)
+			case Green:
+				buf.WriteString(green)
+			case Yellow:
+				buf.WriteString(yellow)
+			case Blue:
+				buf.WriteString(blue)
+			case Magenta:
+				buf.WriteString(magenta)
+			case Cyan:
+				buf.WriteString(cyan)
+			default:
+				buf.WriteString(white)
+			}
+			if seg.Bold {
+				buf.WriteString(bold)
+			}
+		}
+
+		numFinished += seg.Count
+		statusFrac := float64(seg.Count) / float64(s.Total)
+		fNumBlocks := float64(width) * statusFrac
+		numBlocks := int(math.Ceil(fNumBlocks))
+		if seg.Transparent {
+			if numBlocks > 0 {
+				fmt.Fprint(buf, strings.Repeat(string("░"), numBlocks))
+			}
+		} else {
+			if numBlocks > 1 {
+				fmt.Fprint(buf, strings.Repeat(string("▉"), numBlocks-1))
+			}
+			if numBlocks > 0 {
+				frac := fNumBlocks - math.Floor(fNumBlocks)
+				symbol := blockSymbols[int(math.Round(frac*float64(len(blockSymbols)-1)))]
+				fmt.Fprint(buf, string(symbol))
+			}
+		}
+		numBlocksPrinted += numBlocks
+	}
+
+	if width > numBlocksPrinted {
+		fmt.Fprint(buf, strings.Repeat(string(" "), width-numBlocksPrinted))
+	}
+	fmt.Fprint(buf, ansiReset)
+	fmt.Fprint(buf, "] ", percentage(numFinished, s.Total))
+
+	if ansiColors {
+		// move cursor to start of line so the bar is overridden next print
+		fmt.Fprint(buf, positionLeft)
+	} else {
+		// cannot move cursor, so newline
+		fmt.Fprintln(buf)
+	}
+
+	out.Write(buf.Bytes())
+}
+
+// percentage returns the percentage of n out of total as a string
+func percentage(n, total int) string {
+	if total == 0 {
+		return "-"
+	}
+	f := float64(n) / float64(total)
+	return fmt.Sprintf("%.1f%c", f*100.0, '%')
+}