summaryrefslogtreecommitdiffstats
path: root/progress.go
blob: 2d78cef31b9e1ccca78a0a641f682303e83ac181 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
// Package progress implements a simple command line progress bar.
//
// Example:
//
//	import (
//		time"
//
//		"gitlab.com/rhogenson/progress-bar"
//	)
//
//	func main() {
//		b := progress.New(100)
//		for i := range 100 {
//			b.Set(i)
//			b.Print()
//			time.Sleep(time.Second)
//		}
//	}
package progress

import (
	"fmt"
	"os"
	"strings"
	"time"

	"gitlab.com/rhogenson/vecdeque"
	"golang.org/x/term"
)

type measurement struct {
	t time.Time
	i int
}

// Bar is a progess bar.
type Bar struct {
	max          int
	measurements vecdeque.DQ[measurement]
	cols         int
	secondPrint  bool
}

func New(max int) *Bar {
	return &Bar{max: max}
}

// Set sets the current value to val.
func (b *Bar) Set(val int) {
	b.measurements.PushBack(measurement{time.Now(), val})
}

// Print shows the progress bar on standard error.
func (b *Bar) Print() {
	if b.measurements.Len() == 0 {
		return
	}

	now := time.Now()
	for b.measurements.Len() > 1 && now.Sub(b.measurements.Get(0).t) > 2*time.Minute {
		b.measurements.PopFront()
	}
	first := b.measurements.Get(0)
	last := b.measurements.Get(b.measurements.Len() - 1)
	deltaT := last.t.Sub(first.t)
	delta := last.i - first.i
	eta := time.Duration(-1)
	if delta != 0 {
		eta = time.Duration(float64(b.max-last.i) / float64(delta) * float64(deltaT))
	}

	if b.secondPrint {
		fmt.Fprint(os.Stderr, "\033[2F\033[J")
	}
	b.secondPrint = true
	if b.cols == 0 {
		var err error
		if b.cols, _, err = term.GetSize(int(os.Stderr.Fd())); err != nil {
			fmt.Fprintf(os.Stderr, "Unable to determine terminal size: %s\n", err)
		}
	}
	p := last.i * b.cols / b.max
	if b.cols > 0 {
		fmt.Fprintf(os.Stderr, "%s>\n", strings.Repeat("=", max(p-1, 0)))
	}
	if eta < 0 {
		fmt.Fprintln(os.Stderr, "ETA: calculating...")
	} else {
		fmt.Fprintf(os.Stderr, "ETA: %s\n", eta.Round(time.Second))
	}
}