forked from coreos/ioprogress
-
Notifications
You must be signed in to change notification settings - Fork 0
/
progressbar.go
182 lines (159 loc) · 5.14 KB
/
progressbar.go
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
package ioprogress
import (
"fmt"
"io"
"strings"
"sync"
)
var (
// ErrorProgressOutOfBounds is returned if the progress is set to a value
// not between 0 and 1.
ErrorProgressOutOfBounds = fmt.Errorf("progress is out of bounds (0 to 1)")
// ErrorNoBarsAdded is returned when no progress bars have been added to a
// ProgressBarPrinter before PrintAndWait is called.
ErrorNoBarsAdded = fmt.Errorf("AddProgressBar hasn't been called yet")
)
// ProgressBar represents one progress bar in a ProgressBarPrinter. Should not
// be created directly, use the AddProgressBar on a ProgressBarPrinter to
// create these.
type ProgressBar struct {
currentProgress float64
printBefore string
printAfter string
lock sync.Mutex
done bool
}
// SetCurrentProgress sets the progress of this ProgressBar. The progress must
// be between 0 and 1 inclusive.
func (pb *ProgressBar) SetCurrentProgress(progress float64) error {
if progress < 0 || progress > 1 {
return ErrorProgressOutOfBounds
}
pb.lock.Lock()
pb.currentProgress = progress
pb.lock.Unlock()
return nil
}
// SetPrintBefore sets the text printed on the line before the progress bar.
func (pb *ProgressBar) SetPrintBefore(before string) {
pb.lock.Lock()
pb.printBefore = before
pb.lock.Unlock()
}
// SetPrintAfter sets the text printed on the line after the progress bar.
func (pb *ProgressBar) SetPrintAfter(after string) {
pb.lock.Lock()
pb.printAfter = after
pb.lock.Unlock()
}
// ProgressBarPrinter will print out the progress of some number of
// ProgressBars.
type ProgressBarPrinter struct {
// DisplayWidth can be set to influence how large the progress bars are.
// The bars will be scaled to attempt to produce lines of this number of
// characters, but lines of different lengths may still be printed. When
// this value is 0 (aka unset), 80 character columns are assumed.
DisplayWidth int
// PadToBeEven, when set to true, will make Print pad the printBefore text
// with trailing spaces and the printAfter text with leading spaces to make
// the progress bars the same length.
PadToBeEven bool
numLinesInLastPrint int
progressBars []*ProgressBar
lock sync.Mutex
}
// AddProgressBar will create a new ProgressBar, register it with this
// ProgressBarPrinter, and return it. This must be called at least once before
// PrintAndWait is called.
func (pbp *ProgressBarPrinter) AddProgressBar() *ProgressBar {
pb := &ProgressBar{}
pbp.lock.Lock()
pbp.progressBars = append(pbp.progressBars, pb)
pbp.lock.Unlock()
return pb
}
// Print will print out progress information for each ProgressBar that has been
// added to this ProgressBarPrinter. The progress will be written to printTo,
// and if printTo is a terminal it will draw progress bars. AddProgressBar
// must be called at least once before Print is called. If printing to a
// terminal, all draws after the first one will move the cursor up to draw over
// the previously printed bars.
func (pbp *ProgressBarPrinter) Print(printTo io.Writer) (bool, error) {
pbp.lock.Lock()
bars := pbp.progressBars
numColumns := pbp.DisplayWidth
pbp.lock.Unlock()
if len(bars) == 0 {
return false, ErrorNoBarsAdded
}
if numColumns == 0 {
numColumns = 80
}
if isTerminal(printTo) {
moveCursorUp(printTo, pbp.numLinesInLastPrint)
}
maxBefore := 0
maxAfter := 0
for _, bar := range bars {
bar.lock.Lock()
beforeSize := len(bar.printBefore)
afterSize := len(bar.printAfter)
bar.lock.Unlock()
if beforeSize > maxBefore {
maxBefore = beforeSize
}
if afterSize > maxAfter {
maxAfter = afterSize
}
}
allDone := false
for _, bar := range bars {
if isTerminal(printTo) {
bar.printToTerminal(printTo, numColumns, pbp.PadToBeEven, maxBefore, maxAfter)
} else {
bar.printToNonTerminal(printTo)
}
allDone = allDone || bar.currentProgress == 1
}
pbp.numLinesInLastPrint = len(bars)
return allDone, nil
}
// moveCursorUp moves the cursor up numLines in the terminal
func moveCursorUp(printTo io.Writer, numLines int) {
if numLines > 0 {
fmt.Fprintf(printTo, "\033[%dA", numLines)
}
}
func (pb *ProgressBar) printToTerminal(printTo io.Writer, numColumns int, padding bool, maxBefore, maxAfter int) {
pb.lock.Lock()
before := pb.printBefore
after := pb.printAfter
if padding {
before = before + strings.Repeat(" ", maxBefore-len(before))
after = strings.Repeat(" ", maxAfter-len(after)) + after
}
progressBarSize := numColumns - (len(fmt.Sprintf("%s [] %s", before, after)))
progressBar := ""
if progressBarSize > 0 {
currentProgress := int(pb.currentProgress * float64(progressBarSize))
progressBar = fmt.Sprintf("[%s%s] ",
strings.Repeat("=", currentProgress),
strings.Repeat(" ", progressBarSize-currentProgress))
} else {
// If we can't fit the progress bar, better to not pad the before/after.
before = pb.printBefore
after = pb.printAfter
}
fmt.Fprintf(printTo, "%s %s%s\n", before, progressBar, after)
pb.lock.Unlock()
}
func (pb *ProgressBar) printToNonTerminal(printTo io.Writer) {
pb.lock.Lock()
if !pb.done {
fmt.Fprintf(printTo, "%s %s\n", pb.printBefore, pb.printAfter)
if pb.currentProgress == 1 {
pb.done = true
}
}
pb.lock.Unlock()
}