diff --git a/google_drive_downloader.py b/google_drive_downloader.py new file mode 100644 index 0000000..4e3d0cc --- /dev/null +++ b/google_drive_downloader.py @@ -0,0 +1,58 @@ +from __future__ import print_function +import requests +import zipfile +import warnings +from os import makedirs +from os.path import dirname +from os.path import exists + + +class GoogleDriveDownloader: + + CHUNK_SIZE = 32768 + DOWNLOAD_URL = "https://docs.google.com/uc?export=download" + + @staticmethod + def download_file_from_google_drive(file_id, dest_path, overwrite=False, unzip=False): + + destination_directory = dirname(dest_path) + if not exists(destination_directory): + makedirs(destination_directory) + + if not exists(dest_path) or overwrite: + + session = requests.Session() + + print('Downloading {} into {}... '.format(file_id, dest_path), end='', flush=True) + response = session.get(GoogleDriveDownloader.DOWNLOAD_URL, params={'id': file_id}, stream=True) + + token = GoogleDriveDownloader._get_confirm_token(response) + if token: + params = {'id': file_id, 'confirm': token} + response = session.get(GoogleDriveDownloader.DOWNLOAD_URL, params=params, stream=True) + + GoogleDriveDownloader._save_response_content(response, dest_path) + print('Done.') + + if unzip: + try: + print('Unzipping...', end='', flush=True) + with zipfile.ZipFile(dest_path, 'r') as z: + z.extractall(destination_directory) + print('Done.') + except zipfile.BadZipfile: + warnings.warn('Ignoring `unzip` since "{}" does not look like a valid zip file'.format(file_id)) + + @staticmethod + def _get_confirm_token(response): + for key, value in response.cookies.items(): + if key.startswith('download_warning'): + return value + return None + + @staticmethod + def _save_response_content(response, destination): + with open(destination, "wb") as f: + for chunk in response.iter_content(GoogleDriveDownloader.CHUNK_SIZE): + if chunk: # filter out keep-alive new chunks + f.write(chunk)