blob: 11c8f169c7ec95788a488b5cd2ac4113022d5fbe [file] [log] [blame]
// 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, '%')
}