From ba4e05af4f5dd2ab80a8603a823301621a7db55f Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 26 Jul 2020 23:42:46 +0300 Subject: [PATCH] gdrive: upload/download empty file https://github.com/iterative/dvc/issues/4286 --- pydrive2/files.py | 47 ++++++++++++++++++++++++++------------ pydrive2/test/test_file.py | 34 +++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 15 deletions(-) diff --git a/pydrive2/files.py b/pydrive2/files.py index 3065e976..cc77785e 100644 --- a/pydrive2/files.py +++ b/pydrive2/files.py @@ -1,3 +1,4 @@ +import os import io import mimetypes import json @@ -261,7 +262,9 @@ def SetContentString(self, content, encoding="utf-8"): :param content: content of the file in string. :type content: str """ - self.content = io.BytesIO(content.encode(encoding)) + if content: + self.content = io.BytesIO(content.encode(encoding)) + if self.get("mimeType") is None: self["mimeType"] = "text/plain" @@ -275,7 +278,9 @@ def SetContentFile(self, filename): :param filename: name of the file to be uploaded. :type filename: str. """ - self.content = open(filename, "rb") + if os.path.getsize(filename): + self.content = open(filename, "rb") + if self.get("title") is None: self["title"] = filename if self.get("mimeType") is None: @@ -353,20 +358,32 @@ def download(fd, request): download(fd, files.get_media(fileId=file_id)) except errors.HttpError as error: exc = ApiRequestError(error) - if ( - exc.error["code"] != 403 - or exc.GetField("reason") != "fileNotDownloadable" - ): + code = exc.error["code"] + reason = exc.GetField("reason") + if code == 403 and reason == "fileNotDownloadable": + mimetype = mimetype or "text/plain" + fd.seek(0) # just in case `download()` modified `fd` + try: + download( + fd, + files.export_media( + fileId=file_id, mimeType=mimetype + ), + ) + except errors.HttpError as error: + raise ApiRequestError(error) + elif code == 416 and reason == "requestedRangeNotSatisfiable": + # NOTE: An empty file case. Wasting one API call to make + # absolutely sure. See + # https://github.com/iterative/dvc/issues/4507 + try: + self.FetchMetadata(fields="fileSize") + if int(self["fileSize"]) != 0: + raise exc + except errors.HttpError: + raise exc + else: raise exc - mimetype = mimetype or "text/plain" - fd.seek(0) # just in case `download()` modified `fd` - try: - download( - fd, - files.export_media(fileId=file_id, mimeType=mimetype), - ) - except errors.HttpError as error: - raise ApiRequestError(error) if mimetype == "text/plain" and remove_bom: fd.seek(0) diff --git a/pydrive2/test/test_file.py b/pydrive2/test/test_file.py index 4ca0a6a7..a2945c0b 100644 --- a/pydrive2/test/test_file.py +++ b/pydrive2/test/test_file.py @@ -300,6 +300,40 @@ def test_11_Files_Get_Content_Buffer(self): self.DeleteUploadedFiles(drive, [file1["id"]]) + def test_12_Upload_Download_Empty_File(self): + filename = os.path.join(self.tmpdir, str(time())) + create_file(filename, "") + + drive = GoogleDrive(self.ga) + file1 = drive.CreateFile() + file1.SetContentFile(filename) + pydrive_retry(file1.Upload) + + fileOut1 = self.getTempFile() + pydrive_retry(file1.GetContentFile, fileOut1) + self.assertEqual(os.path.getsize(fileOut1), 0) + + self.DeleteUploadedFiles(drive, [file1["id"]]) + + def test_13_Upload_Download_Empty_String(self): + drive = GoogleDrive(self.ga) + file1 = drive.CreateFile() + file1.SetContentString("") + pydrive_retry(file1.Upload) + + self.assertEqual(pydrive_retry(file1.GetContentString), "") + + # Force download and double check content + pydrive_retry(file1.FetchContent) + self.assertEqual(file1.GetContentString(), "") + + # Download file from id + file2 = drive.CreateFile({"id": file1["id"]}) + pydrive_retry(file2.FetchContent) + self.assertEqual(file2.GetContentString(), "") + + self.DeleteUploadedFiles(drive, [file1["id"]]) + # Tests for Trash/UnTrash/Delete. # ===============================