55 lines
1.5 KiB
Go
55 lines
1.5 KiB
Go
package utils
|
|
|
|
import (
|
|
"io"
|
|
"time"
|
|
)
|
|
|
|
// Reader with progress bar
|
|
type progressReader struct {
|
|
reader io.ReadCloser // Stream to read from
|
|
output io.Writer // Where to send progress bar to
|
|
progress JSONProgress
|
|
lastUpdate int // How many bytes read at least update
|
|
ID string
|
|
action string
|
|
sf *StreamFormatter
|
|
newLine bool
|
|
}
|
|
|
|
func (r *progressReader) Read(p []byte) (n int, err error) {
|
|
read, err := r.reader.Read(p)
|
|
r.progress.Current += read
|
|
updateEvery := 1024 * 512 //512kB
|
|
if r.progress.Total > 0 {
|
|
// Update progress for every 1% read if 1% < 512kB
|
|
if increment := int(0.01 * float64(r.progress.Total)); increment < updateEvery {
|
|
updateEvery = increment
|
|
}
|
|
}
|
|
if r.progress.Current-r.lastUpdate > updateEvery || err != nil {
|
|
r.output.Write(r.sf.FormatProgress(r.ID, r.action, &r.progress))
|
|
r.lastUpdate = r.progress.Current
|
|
}
|
|
// Send newline when complete
|
|
if r.newLine && err != nil && read == 0 {
|
|
r.output.Write(r.sf.FormatStatus("", ""))
|
|
}
|
|
return read, err
|
|
}
|
|
func (r *progressReader) Close() error {
|
|
r.progress.Current = r.progress.Total
|
|
r.output.Write(r.sf.FormatProgress(r.ID, r.action, &r.progress))
|
|
return r.reader.Close()
|
|
}
|
|
func ProgressReader(r io.ReadCloser, size int, output io.Writer, sf *StreamFormatter, newline bool, ID, action string) *progressReader {
|
|
return &progressReader{
|
|
reader: r,
|
|
output: NewWriteFlusher(output),
|
|
ID: ID,
|
|
action: action,
|
|
progress: JSONProgress{Total: size, Start: time.Now().UTC().Unix()},
|
|
sf: sf,
|
|
newLine: newline,
|
|
}
|
|
}
|