diff --git a/pydrive2/files.py b/pydrive2/files.py index cd23dfae..bd0b8fa5 100644 --- a/pydrive2/files.py +++ b/pydrive2/files.py @@ -4,6 +4,7 @@ from googleapiclient import errors from googleapiclient.http import MediaIoBaseUpload +from googleapiclient.http import MediaIoBaseDownload from functools import wraps from .apiattr import ApiAttribute @@ -34,6 +35,10 @@ def __init__(self, http_error): # Initialize args for backward compatibility super().__init__(http_error) + def GetField(self, field): + """Returns the `field` from the first error""" + return self.error.get("errors", [{}])[0].get(field, "") + class FileNotDownloadableError(RuntimeError): """Error trying to download file that is not downloadable.""" @@ -220,7 +225,9 @@ def GetContentString( self.FetchContent(mimetype, remove_bom) return self.content.getvalue().decode(encoding) - def GetContentFile(self, filename, mimetype=None, remove_bom=False): + def GetContentFile( + self, filename, mimetype=None, remove_bom=False, callback=None + ): """Save content of this file as a local file. :param filename: name of the file to write to. @@ -229,17 +236,53 @@ def GetContentFile(self, filename, mimetype=None, remove_bom=False): :type mimetype: str :param remove_bom: Whether to remove the byte order marking. :type remove_bom: bool + :param callback: passed two arguments: (total trasferred, file size). + :type param: callable :raises: ApiRequestError, FileNotUploadedError, FileNotDownloadableError """ - if ( - self.content is None - or type(self.content) is not io.BytesIO - or self.has_bom == remove_bom - ): - self.FetchContent(mimetype, remove_bom) - f = open(filename, "wb") - f.write(self.content.getvalue()) - f.close() + files = self.auth.service.files() + file_id = self.metadata.get("id") or self.get("id") + + def download(fd, request): + downloader = MediaIoBaseDownload(fd, request) + done = False + while done is False: + status, done = downloader.next_chunk() + if callback: + callback(status.resumable_progress, status.total_size) + + with open(filename, mode="w+b") as fd: + # Ideally would use files.export_media instead if + # metadata.get("mimeType").startswith("application/vnd.google-apps.") + # but that would first require a slow call to FetchMetadata() + try: + 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" + ): + 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) + boms = [ + bom[mimetype] + for bom in MIME_TYPE_TO_BOM.values() + if mimetype in bom + ] + if boms: + self._RemovePrefix(fd, boms[0]) @LoadAuth def FetchMetadata(self, fields=None, fetch_all=False): diff --git a/pydrive2/test/test_file.py b/pydrive2/test/test_file.py index 6dad4d53..6e49ea47 100644 --- a/pydrive2/test/test_file.py +++ b/pydrive2/test/test_file.py @@ -487,7 +487,7 @@ def test_GFile_Conversion_Lossless_String(self): downloaded_file_name = "_tmp_downloaded_file_name.txt" pydrive_retry( lambda: file1.GetContentFile( - downloaded_file_name, mimetype="text/plain" + downloaded_file_name, mimetype="text/plain", remove_bom=True ) ) downloaded_string = open(downloaded_file_name).read() diff --git a/pydrive2/test/test_util.py b/pydrive2/test/test_util.py index c8143f0e..bba8876d 100644 --- a/pydrive2/test/test_util.py +++ b/pydrive2/test/test_util.py @@ -48,11 +48,7 @@ def pydrive_retry(call): try: result = call() except ApiRequestError as exception: - retry_codes = ["403", "500", "502", "503", "504"] - if any( - "HttpError {}".format(code) in str(exception) - for code in retry_codes - ): + if exception.error["code"] in [403, 500, 502, 503, 504]: raise PyDriveRetriableError("Google API request failed") raise return result