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
|
package ioprogress
import (
"io"
"time"
)
// Reader is an implementation of io.Reader that draws the progress of
// reading some data.
type Reader struct {
// Reader is the underlying reader to read from
Reader io.Reader
// Size is the total size of the data coming out of the reader.
Size int64
// DrawFunc is the callback to invoke to draw the progress bar. By
// default, this will be DrawTerminal(os.Stdout).
//
// DrawInterval is the minimum time to wait between reads to update the
// progress bar.
DrawFunc DrawFunc
DrawInterval time.Duration
progress int64
lastDraw time.Time
}
// Read reads from the underlying reader and invokes the DrawFunc if
// appropriate. The DrawFunc is executed when there is data that is
// read (progress is made) and at least DrawInterval time has passed.
func (r *Reader) Read(p []byte) (int, error) {
// If we haven't drawn before, initialize the progress bar
if r.lastDraw.IsZero() {
r.initProgress()
}
// Read from the underlying source
n, err := r.Reader.Read(p)
// Always increment the progress even if there was an error
r.progress += int64(n)
// If we don't have any errors, then draw the progress. If we are
// at the end of the data, then finish the progress.
if err == nil {
// Only draw if we read data or we've never read data before (to
// initialize the progress bar).
if n > 0 {
r.drawProgress()
}
}
if err == io.EOF {
r.finishProgress()
}
return n, err
}
func (r *Reader) drawProgress() {
// If we've drawn before, then make sure that the draw interval
// has passed before we draw again.
interval := r.DrawInterval
if interval == 0 {
interval = time.Second
}
if !r.lastDraw.IsZero() {
nextDraw := r.lastDraw.Add(interval)
if time.Now().Before(nextDraw) {
return
}
}
// Draw
f := r.drawFunc()
f(r.progress, r.Size)
// Record this draw so that we don't draw again really quickly
r.lastDraw = time.Now()
}
func (r *Reader) finishProgress() {
f := r.drawFunc()
f(r.progress, r.Size)
// Print a newline
f(-1, -1)
// Reset lastDraw so we don't finish again
var zeroDraw time.Time
r.lastDraw = zeroDraw
}
func (r *Reader) initProgress() {
var zeroDraw time.Time
r.lastDraw = zeroDraw
r.drawProgress()
r.lastDraw = zeroDraw
}
func (r *Reader) drawFunc() DrawFunc {
if r.DrawFunc == nil {
return defaultDrawFunc
}
return r.DrawFunc
}
|