Show available formats for given file. Change export arg to --format

4 files changed, 70 insertions, 51 deletions

diff --git a/cli/cli.go b/cli/cli.go
index d713f41..38de1ba 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -331,14 +331,7 @@ func Download(d *gdrive.Drive, fileId string, stdout, deleteAfterDownload bool,
 		return fmt.Errorf("An error occurred: %v\n", err)
 	}
 
-	downloadUrl := info.DownloadUrl
-	extension := ""
-
-	if downloadUrl == "" && format == "" {
-		return fmt.Errorf("Document has no download url. Try to specify an export format.")
-	}
-
-	downloadUrl, extension, err = util.ExportFormat(info, format)
+	downloadUrl, extension, err := util.InternalDownloadUrlAndExtension(info, format)
 	if err != nil {
 		return err
 	}
diff --git a/drive.go b/drive.go
index 680d8c0..e11f7b2 100644
--- a/drive.go
+++ b/drive.go
@@ -53,7 +53,7 @@ type Options struct {
 
 	Download struct {
 		FileId string `goptions:"-i, --id, mutexgroup='download', obligatory, description='File Id'"`
-		Export string `goptions:"-e, --export, description='Export file format for google docs'"`
+		Format string `goptions:"--format, description='Download file in a specified format (needed for google docs)'"`
 		Stdout bool   `goptions:"-s, --stdout, description='Write file content to stdout'"`
 		Force  bool   `goptions:"--force, description='Overwrite existing file'"`
 		Pop    bool   `goptions:"--pop, mutexgroup='download', description='Download latest file, and remove it from google drive'"`
@@ -123,9 +123,9 @@ func main() {
 	case "download":
 		args := opts.Download
 		if args.Pop {
-			err = cli.DownloadLatest(drive, args.Stdout, args.Export, args.Force)
+			err = cli.DownloadLatest(drive, args.Stdout, args.Format, args.Force)
 		} else {
-			err = cli.Download(drive, args.FileId, args.Stdout, false, args.Export, args.Force)
+			err = cli.Download(drive, args.FileId, args.Stdout, false, args.Format, args.Force)
 		}
 
 	case "delete":
diff --git a/util/drive.go b/util/drive.go
index c4ba137..259912f 100644
--- a/util/drive.go
+++ b/util/drive.go
@@ -25,44 +25,60 @@ func ParentList(parents []*drive.ParentReference) string {
 	return strings.Join(ids, ", ")
 }
 
-func ExportFormat(info *drive.File, format string) (downloadUrl string, extension string, err error) {
-	// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
-	switch format {
-	case "docx":
-		extension = ".docx"
-		downloadUrl = info.ExportLinks["application/vnd.openxmlformats-officedocument.wordprocessingml.document"]
-	case "xlsx":
-		extension = ".xlsx"
-		downloadUrl = info.ExportLinks["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"]
-	case "pptx":
-		extension = ".pptx"
-		downloadUrl = info.ExportLinks["application/application/vnd.openxmlformats-officedocument.presentationml.presentation"]
-	case "odf":
-		extension = ".odf"
-		downloadUrl = info.ExportLinks["application/vnd.oasis.opendocument.text"]
-	case "ods":
-		extension = ".ods"
-		downloadUrl = info.ExportLinks["application/x-vnd.oasis.opendocument.spreadsheet"]
-	case "pdf":
-		extension = ".pdf"
-		downloadUrl = info.ExportLinks["application/pdf"]
-	case "rtf":
-		extension = ".rtf"
-		downloadUrl = info.ExportLinks["application/rtf"]
-	case "csv":
-		extension = ".csv"
-		downloadUrl = info.ExportLinks["text/csv"]
-	case "html":
-		extension = ".html"
-		downloadUrl = info.ExportLinks["text/html"]
-	case "txt":
-		extension = ".txt"
-		downloadUrl = info.ExportLinks["text/plain"]
-	case "json":
-		extension = ".json"
-		downloadUrl = info.ExportLinks["application/vnd.google-apps.script+json"]
-	default:
-		err = fmt.Errorf("Unknown export format: %s. Known formats: docx, xlsx, pptx, odf, ods, pdf, rtf, csv, txt, html, json", format)
-	}
-	return
+func InternalDownloadUrlAndExtension(info *drive.File, format string) (downloadUrl string, extension string, err error) {
+    // Make a list of available mime types for this file
+    availableMimeTypes := make([]string, 0)
+    for mime, _ := range info.ExportLinks {
+        availableMimeTypes = append(availableMimeTypes, mime)
+    }
+
+    mimeExtensions := map[string]string{
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
+        "application/application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
+        "application/vnd.oasis.opendocument.text": "odf",
+        "application/x-vnd.oasis.opendocument.spreadsheet": "ods",
+        "application/pdf": "pdf",
+        "application/rtf": "rtf",
+        "text/csv": "csv",
+        "text/html": "html",
+        "text/plain": "txt",
+        "application/vnd.google-apps.script+json": "json",
+    }
+
+    // Make a list of available formats for this file
+    availableFormats :=  make([]string, 0)
+    for _, mime := range availableMimeTypes {
+        if ext, ok := mimeExtensions[mime]; ok {
+            availableFormats = append(availableFormats, ext)
+        }
+    }
+
+    // Return DownloadUrl if no format is specified
+    if format == "" {
+        if info.DownloadUrl == "" {
+            return "", "", fmt.Errorf("A format needs to be specified to download this file (--format). Available formats: %s", strings.Join(availableFormats, ", "))
+        }
+        return info.DownloadUrl, "", nil
+    }
+
+    // Ensure that the specified format is available
+    if !inArray(format, availableFormats) {
+        if len(availableFormats) > 0 {
+            return "", "", fmt.Errorf("Invalid format. Available formats: %s", strings.Join(availableFormats, ", "))
+        } else {
+            return "", "", fmt.Errorf("No export formats are available for this file")
+        }
+    }
+
+    // Grab download url
+    for mime, f := range mimeExtensions {
+        if f == format {
+            downloadUrl = info.ExportLinks[mime]
+            break
+        }
+    }
+
+    extension = "." + format
+    return
 }
diff --git a/util/generic.go b/util/generic.go
index 2b26d60..a4b92be 100644
--- a/util/generic.go
+++ b/util/generic.go
@@ -197,3 +197,13 @@ func columnPadder(items []map[string]string, key string, spacing int) func(strin
 		return column
 	}
 }
+
+func inArray(needle string, haystack []string) bool {
+    for _, x := range haystack {
+        if needle == x {
+            return true
+        }
+    }
+
+    return false
+}