diff --git a/pb.go b/pb.go index 58e6b7d..f92e02b 100644 --- a/pb.go +++ b/pb.go @@ -242,6 +242,7 @@ func (pb *ProgressBar) Add(value int) *ProgressBar { return pb.Add64(int64(value)) } +// Increment atomically increments the progress func (pb *ProgressBar) Increment() *ProgressBar { return pb.Add64(1) } @@ -360,6 +361,7 @@ func (pb *ProgressBar) Finish() *ProgressBar { return pb } +// IsStarted indicates progress bar state func (pb *ProgressBar) IsStarted() bool { pb.mu.RLock() defer pb.mu.RUnlock() @@ -379,6 +381,14 @@ func (pb *ProgressBar) SetTemplate(tmpl ProgressBarTemplate) *ProgressBar { return pb.SetTemplateString(string(tmpl)) } +// NewProxyReader creates a wrapper for given reader, but with progress handle +// Takes io.Reader or io.ReadCloser +// Also, it automatically switches progress bar to handle units as bytes +func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader { + pb.Set(Bytes, true) + return &Reader{r, pb} +} + func (pb *ProgressBar) render() (result string, width int) { defer func() { if r := recover(); r != nil { diff --git a/reader.go b/reader.go new file mode 100644 index 0000000..f2e60a0 --- /dev/null +++ b/reader.go @@ -0,0 +1,26 @@ +package pb + +import ( + "io" +) + +// Reader it's a wrapper for given reader, but with progress handle +type Reader struct { + io.Reader + bar *ProgressBar +} + +// Read reads bytes from wrapped reader and add amount of bytes to progress bar +func (r *Reader) Read(p []byte) (n int, err error) { + n, err = r.Reader.Read(p) + r.bar.Add(n) + return +} + +// Close the wrapped reader when it implements io.Closer +func (r *Reader) Close() (err error) { + if closer, ok := r.Reader.(io.Closer); ok { + return closer.Close() + } + return +} diff --git a/reader_test.go b/reader_test.go new file mode 100644 index 0000000..334bb54 --- /dev/null +++ b/reader_test.go @@ -0,0 +1,58 @@ +package pb + +import ( + "testing" +) + +func TestPBProxyReader(t *testing.T) { + bar := new(ProgressBar) + if bar.GetBool(Bytes) { + t.Errorf("By default bytes must be false") + } + + testReader := new(testReaderCloser) + proxyReader := bar.NewProxyReader(testReader) + + if !bar.GetBool(Bytes) { + t.Errorf("Bytes must be true after call NewProxyReader") + } + + for i := 0; i < 10; i++ { + buf := make([]byte, 10) + n, e := proxyReader.Read(buf) + if e != nil { + t.Errorf("Proxy reader return err: %v", e) + } + if n != len(buf) { + t.Errorf("Proxy reader return unexpected N: %d (wand %d)", n, len(buf)) + } + for _, b := range buf { + if b != 'f' { + t.Errorf("Unexpected read value: %v (want %v)", b, 'f') + } + } + if want := int64((i + 1) * len(buf)); bar.Current() != want { + t.Errorf("Unexpected bar current value: %d (want %d)", bar.Current(), want) + } + } + proxyReader.Close() + if !testReader.closed { + t.Errorf("Reader must be closed after call ProxyReader.Close") + } +} + +type testReaderCloser struct { + closed bool +} + +func (tr *testReaderCloser) Read(p []byte) (n int, err error) { + for i := range p { + p[i] = 'f' + } + return len(p), nil +} + +func (tr *testReaderCloser) Close() (err error) { + tr.closed = true + return +}