diff --git a/bar.go b/bar.go index 659abe5..49ce44d 100644 --- a/bar.go +++ b/bar.go @@ -2,6 +2,7 @@ import ( "fmt" + "io" "strconv" "sync" "time" @@ -125,6 +126,10 @@ func (b *Bar) SetEtaAlpha(a float64) *Bar { b.alpha = a return b +} + +func (b *Bar) ProxyReader(r io.Reader) *Reader { + return &Reader{r, b} } // String returns the string representation of the bar diff --git a/proxyreader.go b/proxyreader.go new file mode 100644 index 0000000..b710ec8 --- /dev/null +++ b/proxyreader.go @@ -0,0 +1,23 @@ +package mpb + +import "io" + +// Reader is io.Reader wrapper, for proxy read bytes +type Reader struct { + io.Reader + bar *Bar +} + +func (r *Reader) Read(p []byte) (int, error) { + n, err := r.Reader.Read(p) + r.bar.Incr(n) + return n, err +} + +// Close the reader when it implements io.Closer +func (r *Reader) Close() error { + if closer, ok := r.Reader.(io.Closer); ok { + return closer.Close() + } + return nil +}