Add quota flag

3 files changed, 36 insertions, 13 deletions

diff --git a/cli/cli.go b/cli/cli.go
index 93bca81..49ba1e7 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -49,7 +49,7 @@ func List(d *gdrive.Drive, query, titleFilter string, maxResults int, sharedStat
 		items = append(items, map[string]string{
 			"Id":      f.Id,
 			"Title":   util.TruncateString(f.Title, 40),
-			"Size":    util.FileSizeFormat(f.FileSize),
+			"Size":    util.FileSizeFormat(f.FileSize, false),
 			"Created": util.ISODateToLocal(f.CreatedDate),
 		})
 	}
@@ -101,16 +101,16 @@ func Info(d *gdrive.Drive, fileId string) error {
 	if err != nil {
 		return fmt.Errorf("An error occurred: %v\n", err)
 	}
-	printInfo(d, info)
+	printInfo(d, info, false)
 	return nil
 }
 
-func printInfo(d *gdrive.Drive, f *drive.File) {
+func printInfo(d *gdrive.Drive, f *drive.File, sizeInBytes bool) {
 	fields := map[string]string{
 		"Id":          f.Id,
 		"Title":       f.Title,
 		"Description": f.Description,
-		"Size":        util.FileSizeFormat(f.FileSize),
+		"Size":        util.FileSizeFormat(f.FileSize, sizeInBytes),
 		"Created":     util.ISODateToLocal(f.CreatedDate),
 		"Modified":    util.ISODateToLocal(f.ModifiedDate),
 		"Owner":       strings.Join(f.OwnerNames, ", "),
@@ -140,7 +140,7 @@ func Folder(d *gdrive.Drive, title string, parentId string, share bool) error {
 	if err != nil {
 		return err
 	}
-	printInfo(d, info)
+	printInfo(d, info, false)
 	fmt.Printf("Folder '%s' created\n", info.Title)
 	return nil
 }
@@ -189,9 +189,9 @@ func UploadStdin(d *gdrive.Drive, input io.ReadCloser, title string, parentId st
 	bytes := info.FileSize
 
 	// Print information about uploaded file
-	printInfo(d, info)
+	printInfo(d, info, false)
 	fmt.Printf("MIME Type: %s\n", mimeType)
-	fmt.Printf("Uploaded '%s' at %s, total %s\n", info.Title, getRate(bytes), util.FileSizeFormat(bytes))
+	fmt.Printf("Uploaded '%s' at %s, total %s\n", info.Title, getRate(bytes), util.FileSizeFormat(bytes, false))
 
 	// Share file if the share flag was provided
 	if share {
@@ -297,9 +297,9 @@ func uploadFile(d *gdrive.Drive, input *os.File, inputInfo os.FileInfo, title st
 	bytes := info.FileSize
 
 	// Print information about uploaded file
-	printInfo(d, info)
+	printInfo(d, info, false)
 	fmt.Printf("MIME Type: %s\n", mimeType)
-	fmt.Printf("Uploaded '%s' at %s, total %s\n", info.Title, getRate(bytes), util.FileSizeFormat(bytes))
+	fmt.Printf("Uploaded '%s' at %s, total %s\n", info.Title, getRate(bytes), util.FileSizeFormat(bytes, false))
 
 	// Share file if the share flag was provided
 	if share {
@@ -375,7 +375,7 @@ func Download(d *gdrive.Drive, fileId string, stdout, deleteAfterDownload bool,
 		return fmt.Errorf("An error occurred: %s", err)
 	}
 
-	fmt.Printf("Downloaded '%s' at %s, total %s\n", fileName, getRate(bytes), util.FileSizeFormat(bytes))
+	fmt.Printf("Downloaded '%s' at %s, total %s\n", fileName, getRate(bytes), util.FileSizeFormat(bytes, false))
 
 	if deleteAfterDownload {
 		err = Delete(d, fileId)
@@ -435,6 +435,18 @@ func Unshare(d *gdrive.Drive, fileId string) error {
 	return nil
 }
 
+func Quota(d *gdrive.Drive, sizeInBytes bool) error {
+	info, err := d.About.Get().Do()
+	if err != nil {
+		return fmt.Errorf("An error occurred: %v\n", err)
+	}
+
+	fmt.Printf("Used: %s\n", util.FileSizeFormat(info.QuotaBytesUsed, sizeInBytes))
+	fmt.Printf("Free: %s\n", util.FileSizeFormat(info.QuotaBytesTotal-info.QuotaBytesUsed, sizeInBytes))
+	fmt.Printf("Total: %s\n", util.FileSizeFormat(info.QuotaBytesTotal, sizeInBytes))
+	return nil
+}
+
 func isShared(d *gdrive.Drive, fileId string) bool {
 	r, err := d.Permissions.List(fileId).Do()
 	if err != nil {
diff --git a/drive.go b/drive.go
index bc4bfd4..1786e68 100644
--- a/drive.go
+++ b/drive.go
@@ -77,6 +77,10 @@ type Options struct {
 		Preview  bool   `goptions:"-p, --preview, mutexgroup='urltype', description='Generate preview url (default)'"`
 		Download bool   `goptions:"-d, --download, mutexgroup='urltype', description='Generate download url'"`
 	} `goptions:"url"`
+
+	Quota struct {
+		SizeInBytes bool `goptions:"--bytes, description='Show size in bytes'"`
+	} `goptions:"quota"`
 }
 
 func main() {
@@ -145,6 +149,9 @@ func main() {
 			fmt.Println(util.PreviewUrl(opts.Url.FileId))
 		}
 
+	case "quota":
+		err = cli.Quota(drive, opts.Quota.SizeInBytes)
+
 	default:
 		goptions.PrintHelp()
 	}
diff --git a/util/generic.go b/util/generic.go
index 50c22f5..ded985d 100644
--- a/util/generic.go
+++ b/util/generic.go
@@ -48,7 +48,11 @@ func FormatBool(b bool) string {
 	return strings.Title(strconv.FormatBool(b))
 }
 
-func FileSizeFormat(bytes int64) string {
+func FileSizeFormat(bytes int64, forceBytes bool) string {
+	if forceBytes {
+		return fmt.Sprintf("%v B", bytes)
+	}
+
 	units := []string{"B", "KB", "MB", "GB", "TB", "PB"}
 
 	var i int
@@ -122,10 +126,10 @@ func MeasureTransferRate() func(int64) string {
 	return func(bytes int64) string {
 		seconds := int64(time.Now().Sub(start).Seconds())
 		if seconds < 1 {
-			return fmt.Sprintf("%s/s", FileSizeFormat(bytes))
+			return fmt.Sprintf("%s/s", FileSizeFormat(bytes, false))
 		}
 		bps := bytes / seconds
-		return fmt.Sprintf("%s/s", FileSizeFormat(bps))
+		return fmt.Sprintf("%s/s", FileSizeFormat(bps, false))
 	}
 }