import os import re import six import sys import json import tqdm import time import shutil import warnings import tempfile import textwrap import requests from six.moves import urllib_parse def indent(text, prefix): """Indent each non-empty line of text with the given prefix.""" return "".join( (prefix + line if line.strip() else line) for line in text.splitlines(True) ) class FileURLRetrievalError(Exception): pass class FolderContentsMaximumLimitError(Exception): pass def parse_url(url, warning=True): """Parse URLs especially for Google Drive links. Args: url: URL to parse. warning: Whether to warn if the URL is not a download link. Returns: A tuple (file_id, is_download_link), where file_id is the ID of the file on Google Drive, and is_download_link is a flag indicating whether the URL is a download link. """ parsed = urllib_parse.urlparse(url) query = urllib_parse.parse_qs(parsed.query) is_gdrive = parsed.hostname in ("drive.google.com", "docs.google.com") is_download_link = parsed.path.endswith("/uc") if not is_gdrive: return None, is_download_link file_id = query.get("id", [None])[0] if file_id is None: for pattern in ( r"^/file/d/(.*?)/(edit|view)$", r"^/file/u/[0-9]+/d/(.*?)/(edit|view)$", r"^/document/d/(.*?)/(edit|htmlview|view)$", r"^/document/u/[0-9]+/d/(.*?)/(edit|htmlview|view)$", r"^/presentation/d/(.*?)/(edit|htmlview|view)$", r"^/presentation/u/[0-9]+/d/(.*?)/(edit|htmlview|view)$", r"^/spreadsheets/d/(.*?)/(edit|htmlview|view)$", r"^/spreadsheets/u/[0-9]+/d/(.*?)/(edit|htmlview|view)$", ): match = re.match(pattern, parsed.path) if match: file_id = match.group(1) break if warning and not is_download_link: warnings.warn( "You specified a Google Drive link that is not the correct link " "to download a file. You might want to try `--fuzzy` option " f"or the following url: https://drive.google.com/uc?id={file_id}" ) return file_id, is_download_link CHUNK_SIZE = 512 * 1024 # 512KB HOME = os.path.expanduser("~") def get_url_from_gdrive_confirmation(contents): """Extract the download URL from a Google Drive confirmation page.""" for pattern in ( r'href="(\/uc\?export=download[^"]+)', r'href="/open\?id=([^"]+)"', r'"downloadUrl":"([^"]+)', ): match = re.search(pattern, contents) if match: url = match.group(1) if pattern == r'href="/open\?id=([^"]+)"': uuid = re.search( r'(.*)
', contents) if match: error = match.group(1) raise FileURLRetrievalError(error) raise FileURLRetrievalError( "Cannot retrieve the public link of the file. " "You may need to change the permission to " "'Anyone with the link', or have had many accesses." ) def _get_session(proxy, use_cookies, return_cookies_file=False): """Create a requests session with optional proxy and cookie handling.""" sess = requests.session() sess.headers.update( {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6)"} ) if proxy is not None: sess.proxies = {"http": proxy, "https": proxy} print("Using proxy:", proxy, file=sys.stderr) cookies_file = os.path.join(HOME, ".cache/gdown/cookies.json") if os.path.exists(cookies_file) and use_cookies: with open(cookies_file) as f: cookies = json.load(f) for k, v in cookies: sess.cookies[k] = v return (sess, cookies_file) if return_cookies_file else sess def download( url=None, output=None, quiet=False, proxy=None, speed=None, use_cookies=True, verify=True, id=None, fuzzy=True, resume=False, format=None, ): """Download file from URL. Parameters ---------- url: str URL. Google Drive URL is also supported. output: str Output filename. Default is basename of URL. quiet: bool Suppress terminal output. Default is False. proxy: str Proxy. speed: float Download byte size per second (e.g., 256KB/s = 256 * 1024). use_cookies: bool Flag to use cookies. Default is True. verify: bool or string Either a bool, in which case it controls whether the server's TLS certificate is verified, or a string, in which case it must be a path to a CA bundle to use. Default is True. id: str Google Drive's file ID. fuzzy: bool Fuzzy extraction of Google Drive's file Id. Default is False. resume: bool Resume the download from existing tmp file if possible. Default is False. format: str, optional Format of Google Docs, Spreadsheets and Slides. Default is: - Google Docs: 'docx' - Google Spreadsheet: 'xlsx' - Google Slides: 'pptx' Returns ------- output: str Output filename. """ if not (id is None) ^ (url is None): raise ValueError("Either url or id has to be specified") if id is not None: url = f"https://drive.google.com/uc?id={id}" url_origin = url sess, cookies_file = _get_session( proxy=proxy, use_cookies=use_cookies, return_cookies_file=True ) gdrive_file_id, is_gdrive_download_link = parse_url(url, warning=not fuzzy) if fuzzy and gdrive_file_id: # overwrite the url with fuzzy match of a file id url = f"https://drive.google.com/uc?id={gdrive_file_id}" url_origin = url is_gdrive_download_link = True while True: res = sess.get(url, stream=True, verify=verify) if url == url_origin and res.status_code == 500: # The file could be Google Docs or Spreadsheets. url = f"https://drive.google.com/open?id={gdrive_file_id}" continue if res.headers["Content-Type"].startswith("text/html"): title = re.search("