Change output text, refactoring

Foxsky28 · Feb 14, 2016 · 453c097 · 453c097
1 parent 2b11e0b
commit 453c097
Show file tree

Hide file tree

Showing 9 changed files with 103 additions and 94 deletions.
diff --git a/drive/download.go b/drive/download.go
@@ -3,7 +3,6 @@ package drive
 import (
     "fmt"
     "io"
-    "io/ioutil"
     "os"
     "time"
     "path/filepath"
@@ -22,49 +21,64 @@ type DownloadArgs struct {
 }
 
 func (self *Drive) Download(args DownloadArgs) error {
-    return self.download(args)
-}
+    if args.Recursive {
+        return self.downloadRecursive(args)
+    }
 
-func (self *Drive) download(args DownloadArgs) error {
     f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
     if err != nil {
         return fmt.Errorf("Failed to get file: %s", err)
     }
 
-    if isDir(f) && !args.Recursive {
+    if isDir(f) {
         return fmt.Errorf("'%s' is a directory, use --recursive to download directories", f.Name)
-    } else if isDir(f) && args.Recursive {
+    }
+
+    if !isBinary(f) {
+        return fmt.Errorf("'%s' is a google document and must be exported, see the export command", f.Name)
+    }
+
+    bytes, rate, err := self.downloadBinary(f, args)
+
+    if !args.Stdout {
+        fmt.Fprintf(args.Out, "Downloaded %s at %s/s, total %s\n", f.Id, formatSize(rate, false), formatSize(bytes, false))
+    }
+    return err
+}
+
+func (self *Drive) downloadRecursive(args DownloadArgs) error {
+    f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
+    if err != nil {
+        return fmt.Errorf("Failed to get file: %s", err)
+    }
+
+    if isDir(f) {
         return self.downloadDirectory(f, args)
     } else if isBinary(f) {
-        return self.downloadBinary(f, args)
-    } else if !args.Recursive {
-        return fmt.Errorf("'%s' is a google document and must be exported, see the export command", f.Name)
+        _, _, err = self.downloadBinary(f, args)
+        return err
     }
 
     return nil
 }
 
-func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) error {
+func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) (int64, int64, error) {
     res, err := self.service.Files.Get(f.Id).Download()
     if err != nil {
-        return fmt.Errorf("Failed to download file: %s", err)
+        return 0, 0, fmt.Errorf("Failed to download file: %s", err)
     }
 
     // Close body on function exit
     defer res.Body.Close()
 
-    // Discard other output if file is written to stdout
-    out := args.Out
-    if args.Stdout {
-        out = ioutil.Discard
-    }
-
     // Path to file
     fpath := filepath.Join(args.Path, f.Name)
 
-    fmt.Fprintf(out, "Downloading %s -> %s\n", f.Name, fpath)
+    if !args.Stdout {
+        fmt.Fprintf(args.Out, "Downloading %s -> %s\n", f.Name, fpath)
+    }
 
-    bytes, rate, err := self.saveFile(saveFileArgs{
+    return self.saveFile(saveFileArgs{
         out: args.Out,
         body: res.Body,
         contentLength: res.ContentLength,
@@ -73,13 +87,6 @@ func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) error {
         stdout: args.Stdout,
         progress: args.Progress,
     })
-
-    if err != nil {
-        return err
-    }
-
-    fmt.Fprintf(out, "Download complete, rate: %s/s, total size: %s\n", formatSize(rate, false), formatSize(bytes, false))
-    return nil
 }
 
 type saveFileArgs struct {
@@ -164,7 +171,7 @@ func (self *Drive) downloadDirectory(parent *drive.File, args DownloadArgs) erro
         newArgs.Id = f.Id
         newArgs.Stdout = false
 
-        err = self.download(newArgs)
+        err = self.downloadRecursive(newArgs)
         if err != nil {
             return err
         }

diff --git a/drive/export.go b/drive/export.go
@@ -24,7 +24,7 @@ type ExportArgs struct {
     Force bool
 }
 
-func (self *Drive) Export(args ExportArgs) (err error) {
+func (self *Drive) Export(args ExportArgs) error {
     f, err := self.service.Files.Get(args.Id).Fields("name", "mimeType").Do()
     if err != nil {
         return fmt.Errorf("Failed to get file: %s", err)
@@ -64,13 +64,13 @@ func (self *Drive) Export(args ExportArgs) (err error) {
     defer outFile.Close()
 
     // Save file to disk
-    bytes, err := io.Copy(outFile, res.Body)
+    _, err = io.Copy(outFile, res.Body)
     if err != nil {
         return fmt.Errorf("Failed saving file: %s", err)
     }
 
-    fmt.Fprintf(args.Out, "Exported '%s' at %s, total %d\n", filename, "x/s", bytes)
-    return
+    fmt.Fprintf(args.Out, "Exported '%s' with mime type: '%s'\n", filename, exportMime)
+    return nil
 }
 
 func (self *Drive) printMimes(out io.Writer, mimeType string) error {
@@ -103,7 +103,7 @@ func getExportMime(userMime, fileMime string) (string, error) {
 
 func getExportFilename(name, mimeType string) string {
     extensions, err := mime.ExtensionsByType(mimeType)
-    if err != nil {
+    if err != nil || len(extensions) == 0 {
         return name
     }
 

diff --git a/drive/import.go b/drive/import.go
@@ -33,7 +33,7 @@ func (self *Drive) Import(args ImportArgs) error {
         return fmt.Errorf("Mime type '%s' is not supported for import", fromMime)
     }
 
-    f, err := self.uploadFile(UploadArgs{
+    f, _, err := self.uploadFile(UploadArgs{
         Out: ioutil.Discard,
         Progress: args.Progress,
         Path: args.Path,
@@ -45,9 +45,7 @@ func (self *Drive) Import(args ImportArgs) error {
         return err
     }
 
-    fmt.Fprintf(args.Out, "[document] id: %s, name: %s\n", f.Id, f.Name)
-    fmt.Fprintf(args.Out, "Imported %s with mime type: '%s'\n", args.Path, toMimes[0])
-
+    fmt.Fprintf(args.Out, "Imported %s with mime type: '%s'\n", f.Id, toMimes[0])
     return nil
 }
 

diff --git a/drive/mkdir.go b/drive/mkdir.go
@@ -20,7 +20,7 @@ func (self *Drive) Mkdir(args MkdirArgs) error {
     if err != nil {
         return err
     }
-    fmt.Printf("Directory '%s' created\n", f.Name)
+    fmt.Fprintf(args.Out, "Directory %s created\n", f.Id)
     return nil
 }
 
@@ -36,8 +36,6 @@ func (self *Drive) mkdir(args MkdirArgs) (*drive.File, error) {
         return nil, fmt.Errorf("Failed to create directory: %s", err)
     }
 
-    fmt.Fprintf(args.Out, "\n[directory] id: %s, name: %s\n", f.Id, f.Name)
-
     //if args.Share {
     //    self.share(TODO)
     //}

diff --git a/drive/update.go b/drive/update.go
@@ -3,7 +3,7 @@ package drive
 import (
     "fmt"
     "mime"
-    "os"
+    "time"
     "io"
     "path/filepath"
     "google.golang.org/api/googleapi"
@@ -19,25 +19,17 @@ type UpdateArgs struct {
     Parents []string
     Mime string
     Recursive bool
-    Stdin bool
     Share bool
     ChunkSize int64
 }
 
-func (self *Drive) Update(args UpdateArgs) (err error) {
-    //if args.Stdin {
-    //    self.uploadStdin()
-    //}
-
-    srcFile, err := os.Open(args.Path)
+func (self *Drive) Update(args UpdateArgs) error {
+    srcFile, srcFileInfo, err := openFile(args.Path)
     if err != nil {
         return fmt.Errorf("Failed to open file: %s", err)
     }
 
-    srcFileInfo, err := srcFile.Stat()
-    if err != nil {
-        return fmt.Errorf("Failed to read file metadata: %s", err)
-    }
+    defer srcFile.Close()
 
     // Instantiate empty drive file
     dstFile := &drive.File{}
@@ -65,14 +57,17 @@ func (self *Drive) Update(args UpdateArgs) (err error) {
     // Wrap file in progress reader
     srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
 
-    f, err := self.service.Files.Update(args.Id, dstFile).Media(srcReader, chunkSize).Do()
+    fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
+    started := time.Now()
+
+    f, err := self.service.Files.Update(args.Id, dstFile).Fields("id", "name", "size").Media(srcReader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }
 
-    fmt.Fprintf(args.Out, "Uploaded '%s' at %s, total %d\n", f.Name, "x/s", f.Size)
-    //if args.Share {
-    //    self.Share(TODO)
-    //}
-    return
+    // Calculate average upload rate
+    rate := calcRate(f.Size, started, time.Now())
+
+    fmt.Fprintf(args.Out, "Updated %s at %s/s, total %s\n", f.Id, formatSize(rate, false), formatSize(f.Size, false))
+    return nil
 }
diff --git a/drive/upload.go b/drive/upload.go
@@ -28,22 +28,35 @@ func (self *Drive) Upload(args UploadArgs) error {
         return fmt.Errorf("Chunk size is to big, max chunk size for this computer is %d", intMax() - 1)
     }
 
-    return self.upload(args)
-}
+    if args.Recursive {
+        return self.uploadRecursive(args)
+    }
 
-func (self *Drive) upload(args UploadArgs) error {
     info, err := os.Stat(args.Path)
     if err != nil {
         return fmt.Errorf("Failed stat file: %s", err)
     }
 
-    if info.IsDir() && !args.Recursive {
+    if info.IsDir() {
         return fmt.Errorf("'%s' is a directory, use --recursive to upload directories", info.Name())
-    } else if info.IsDir() {
+    }
+
+    f, rate, err := self.uploadFile(args)
+    fmt.Fprintf(args.Out, "Uploaded %s at %s/s, total %s\n", f.Id, formatSize(rate, false), formatSize(f.Size, false))
+    return err
+}
+
+func (self *Drive) uploadRecursive(args UploadArgs) error {
+    info, err := os.Stat(args.Path)
+    if err != nil {
+        return fmt.Errorf("Failed stat file: %s", err)
+    }
+
+    if info.IsDir() {
         args.Name = ""
         return self.uploadDirectory(args)
     } else {
-        _, err := self.uploadFile(args)
+        _, _, err := self.uploadFile(args)
         return err
     }
 }
@@ -57,6 +70,7 @@ func (self *Drive) uploadDirectory(args UploadArgs) error {
     // Close file on function exit
     defer srcFile.Close()
 
+    fmt.Fprintf(args.Out, "Creating directory %s\n", srcFileInfo.Name())
     // Make directory on drive
     f, err := self.mkdir(MkdirArgs{
         Out: args.Out,
@@ -81,7 +95,7 @@ func (self *Drive) uploadDirectory(args UploadArgs) error {
         newArgs.Parents = []string{f.Id}
 
         // Upload
-        err = self.upload(newArgs)
+        err = self.uploadRecursive(newArgs)
         if err != nil {
             return err
         }
@@ -90,10 +104,10 @@ func (self *Drive) uploadDirectory(args UploadArgs) error {
     return nil
 }
 
-func (self *Drive) uploadFile(args UploadArgs) (*drive.File, error) {
+func (self *Drive) uploadFile(args UploadArgs) (*drive.File, int64, error) {
     srcFile, srcFileInfo, err := openFile(args.Path)
     if err != nil {
-        return nil, err
+        return nil, 0, err
     }
 
     // Close file on function exit
@@ -125,20 +139,18 @@ func (self *Drive) uploadFile(args UploadArgs) (*drive.File, error) {
     // Wrap file in progress reader
     srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
 
-    fmt.Fprintf(args.Out, "\nUploading %s...\n", args.Path)
+    fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
     started := time.Now()
 
     f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum").Media(srcReader, chunkSize).Do()
     if err != nil {
-        return nil, fmt.Errorf("Failed to upload file: %s", err)
+        return nil, 0, fmt.Errorf("Failed to upload file: %s", err)
     }
 
     // Calculate average upload rate
     rate := calcRate(f.Size, started, time.Now())
 
-    fmt.Fprintf(args.Out, "[file] id: %s, md5: %s, name: %s\n", f.Id, f.Md5Checksum, f.Name)
-    fmt.Fprintf(args.Out, "Uploaded '%s' at %s/s, total %s\n", f.Name, formatSize(rate, false), formatSize(f.Size, false))
-    return f, nil
+    return f, rate, nil
 }
 
 type UploadStreamArgs struct {
@@ -170,28 +182,20 @@ func (self *Drive) UploadStream(args UploadStreamArgs) (err error) {
     // Chunk size option
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
-    f, err := self.service.Files.Create(dstFile).Media(args.In, chunkSize).Do()
+    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()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }
 
-    fmt.Fprintf(args.Out, "Uploaded '%s' at %s, total %d\n", f.Name, "x/s", f.Size)
+    // Calculate average upload rate
+    rate := calcRate(f.Size, started, time.Now())
+
+    fmt.Fprintf(args.Out, "Uploaded %s at %s/s, total %s\n", f.Id, formatSize(rate, false), formatSize(f.Size, false))
     //if args.Share {
     //    self.Share(TODO)
     //}
     return
 }
-
-func openFile(path string) (*os.File, os.FileInfo, error) {
-    f, err := os.Open(path)
-    if err != nil {
-        return nil, nil, fmt.Errorf("Failed to open file: %s", err)
-    }
-
-    info, err := f.Stat()
-    if err != nil {
-        return nil, nil, fmt.Errorf("Failed getting file metadata: %s", err)
-    }
-
-    return f, info, nil
-}
diff --git a/drive/util.go b/drive/util.go
@@ -153,3 +153,17 @@ func min(x int, y int) int {
     n := math.Min(float64(x), float64(y))
     return int(n)
 }
+
+func openFile(path string) (*os.File, os.FileInfo, error) {
+    f, err := os.Open(path)
+    if err != nil {
+        return nil, nil, fmt.Errorf("Failed to open file: %s", err)
+    }
+
+    info, err := f.Stat()
+    if err != nil {
+        return nil, nil, fmt.Errorf("Failed getting file metadata: %s", err)
+    }
+
+    return f, info, nil
+}