diff --git a/bar.go b/bar.go index a643df3..ad81d38 100644 --- a/bar.go +++ b/bar.go @@ -4,6 +4,7 @@ "bytes" "fmt" "io" + "io/ioutil" "strings" "sync" "time" @@ -143,9 +144,16 @@ } } -// ProxyReader allows progress tracking against provided io.Reader. +// ProxyReader wraps r with metrics required for progress tracking. func (b *Bar) ProxyReader(r io.Reader) io.Reader { - return &proxyReader{r, b} + if r == nil { + panic("expect io.Reader, got nil") + } + rc, ok := r.(io.ReadCloser) + if !ok { + rc = ioutil.NopCloser(r) + } + return &proxyReader{rc, b, time.Now()} } // ID returs id of the bar. diff --git a/proxyreader.go b/proxyreader.go index 2bd30f6..d2692cc 100644 --- a/proxyreader.go +++ b/proxyreader.go @@ -7,22 +7,16 @@ // proxyReader is io.Reader wrapper, for proxy read bytes type proxyReader struct { - r io.Reader - b *Bar + io.ReadCloser + bar *Bar + iT time.Time } -func (s *proxyReader) Read(p []byte) (n int, err error) { - start := time.Now() - n, err = s.r.Read(p) +func (pr *proxyReader) Read(p []byte) (n int, err error) { + n, err = pr.ReadCloser.Read(p) if n > 0 { - s.b.IncrBy(n, time.Since(start)) + pr.bar.IncrBy(n, time.Since(pr.iT)) + pr.iT = time.Now() } return } - -func (s *proxyReader) Close() error { - if closer, ok := s.r.(io.Closer); ok { - return closer.Close() - } - return nil -}