Add progress bar to upload from stdin

2 files changed, 12 insertions, 3 deletions

diff --git a/drive/progress.go b/drive/progress.go
index 9d4eb5a..989191e 100644
--- a/drive/progress.go
+++ b/drive/progress.go
@@ -12,7 +12,7 @@ const MaxRateInterval = time.Second * 3
 
 func getProgressReader(r io.Reader, w io.Writer, size int64) io.Reader {
     // Don't wrap reader if output is discarded or size is too small
-    if w == ioutil.Discard || size < 1024 * 1024 {
+    if w == ioutil.Discard || (size > 0 && size < 1024 * 1024) {
         return r
     }
 
@@ -79,7 +79,12 @@ func (self *Progress) draw(isLast bool) {
     self.clear()
 
     // Print progress
-    fmt.Fprintf(self.Writer, "%s/%s", formatSize(self.progress, false), formatSize(self.Size, false))
+    fmt.Fprintf(self.Writer, "%s", formatSize(self.progress, false))
+
+    // Print total size
+    if self.Size > 0 {
+        fmt.Fprintf(self.Writer, "/%s", formatSize(self.Size, false))
+    }
 
     // Print rate
     if self.rate > 0 {
diff --git a/drive/upload.go b/drive/upload.go
index ad3cb34..4f237ca 100644
--- a/drive/upload.go
+++ b/drive/upload.go
@@ -161,6 +161,7 @@ type UploadStreamArgs struct {
     Mime string
     Share bool
     ChunkSize int64
+    Progress io.Writer
 }
 
 func (self *Drive) UploadStream(args UploadStreamArgs) (err error) {
@@ -182,10 +183,13 @@ func (self *Drive) UploadStream(args UploadStreamArgs) (err error) {
     // Chunk size option
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
+    // Wrap file in progress reader
+    srcReader := getProgressReader(args.In, args.Progress, 0)
+
     fmt.Fprintf(args.Out, "Uploading %s\n", dstFile.Name)
     started := time.Now()
 
-    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size").Media(args.In, chunkSize).Do()
+    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size").Media(srcReader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }