#114 add proxy reader

Author: cheggaaa

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 {

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
+}

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
+}