Add progress bar to upload from stdin

wzen-fork · Feb 14, 2016 · ef097ce · ef097ce
1 parent 453c097
commit ef097ce
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 3 deletions.
diff --git 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
@@ -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)
     }

diff --git a/gdrive.go b/gdrive.go
@@ -193,6 +193,12 @@ func main() {
                         Description: "Share file",
                         OmitValue: true,
                     },
+                    cli.BoolFlag{
+                        Name: "noProgress",
+                        Patterns: []string{"--no-progress"},
+                        Description: "Hide progress",
+                        OmitValue: true,
+                    },
                 ),
             },
         },

diff --git a/handlers_drive.go b/handlers_drive.go
@@ -114,6 +114,7 @@ func uploadStdinHandler(ctx cli.Context) {
         Mime: args.String("mime"),
         Share: args.Bool("share"),
         ChunkSize: args.Int64("chunksize"),
+        Progress: progressWriter(args.Bool("noProgress")),
     })
     checkErr(err)
 }