Initial progress indicator

1 files changed, 86 insertions, 0 deletions

diff --git a/drive/util.go b/drive/util.go
index af80b20..366f32b 100644
--- a/drive/util.go
+++ b/drive/util.go
@@ -1,11 +1,13 @@
 package drive
 
 import (
+    "io"
     "os"
     "fmt"
     "strings"
     "strconv"
     "unicode/utf8"
+    "math"
     "time"
 )
 
@@ -39,6 +41,21 @@ func formatSize(bytes int64, forceBytes bool) string {
     return fmt.Sprintf("%.1f %s", value, units[i])
 }
 
+func calcRate(bytes int64, start, end time.Time) int64 {
+    seconds := float64(end.Sub(start).Seconds())
+    if seconds < 1.0 {
+        return bytes
+    }
+    return round(float64(bytes) / seconds)
+}
+
+func round(n float64) int64 {
+    if n < 0 {
+        return int64(math.Ceil(n - 0.5))
+    }
+    return int64(math.Floor(n + 0.5))
+}
+
 func formatBool(b bool) string {
     return strings.Title(strconv.FormatBool(b))
 }
@@ -109,3 +126,72 @@ func fileExists(path string) bool {
 func intMax() int64 {
     return 1 << (strconv.IntSize - 1) - 1
 }
+
+type Progress struct {
+    Writer io.Writer
+    Reader io.Reader
+    Size int64
+    progress int64
+    rate int64
+    rateProgress int64
+    rateUpdated time.Time
+    updated time.Time
+    done bool
+}
+
+func (self *Progress) Read(p []byte) (int, error) {
+    // Read
+    n, err := self.Reader.Read(p)
+
+    now := time.Now()
+    isLast := err != nil
+
+    // Increment progress
+    newProgress := self.progress + int64(n)
+    self.progress = newProgress
+
+    if self.rateUpdated.IsZero() {
+        self.rateUpdated = now
+    }
+
+    // Update rate every 3 seconds
+    if self.rateUpdated.Add(time.Second * 3).Before(now) {
+        self.rate = calcRate(newProgress - self.rateProgress, self.rateUpdated, now)
+        self.rateUpdated = now
+        self.rateProgress = newProgress
+    }
+
+    // Draw progress every second
+    if self.updated.Add(time.Second).Before(now) || isLast {
+        self.Draw(isLast)
+    }
+
+    // Update last draw time
+    self.updated = now
+
+    // Mark as done if error occurs
+    self.done = isLast
+
+    return n, err
+}
+
+func (self *Progress) Draw(isLast bool) {
+    if self.done {
+        return
+    }
+
+    // Clear line
+    fmt.Fprintf(self.Writer, "\r%50s", "")
+
+    // Print progress
+    fmt.Fprintf(self.Writer, "\r%s/%s", formatSize(self.progress, false), formatSize(self.Size, false))
+
+    // Print rate
+    if self.rate > 0 {
+        fmt.Fprintf(self.Writer, ", Rate: %s/s", formatSize(self.rate, false))
+    }
+
+    if isLast {
+        fmt.Fprintf(self.Writer, "\n")
+    }
+}