Spaces:
Running
Running
from __future__ import print_function | |
import json | |
import os | |
import os.path as osp | |
import re | |
import warnings | |
from six.moves import urllib_parse | |
import shutil | |
import sys | |
import tempfile | |
import textwrap | |
import time | |
import requests | |
import six | |
import tqdm | |
def indent(text, prefix): | |
def prefixed_lines(): | |
for line in text.splitlines(True): | |
yield (prefix + line if line.strip() else line) | |
return "".join(prefixed_lines()) | |
class FileURLRetrievalError(Exception): | |
pass | |
class FolderContentsMaximumLimitError(Exception): | |
pass | |
def parse_url(url, warning=True): | |
"""Parse URLs especially for Google Drive links. | |
file_id: ID of file on Google Drive. | |
is_download_link: Flag if it is download link of Google Drive. | |
""" | |
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 is_gdrive, is_download_link | |
file_id = None | |
if "id" in query: | |
file_ids = query["id"] | |
if len(file_ids) == 1: | |
file_id = file_ids[0] | |
else: | |
patterns = [ | |
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)$", | |
] | |
for pattern in patterns: | |
match = re.match(pattern, parsed.path) | |
if match: | |
file_id = match.groups()[0] | |
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 " | |
"or the following url: {url}".format( | |
url="https://drive.google.com/uc?id={}".format(file_id) | |
) | |
) | |
return file_id, is_download_link | |
CHUNK_SIZE = 512 * 1024 # 512KB | |
home = osp.expanduser("~") | |
def get_url_from_gdrive_confirmation(contents): | |
url = "" | |
m = re.search(r'href="(\/uc\?export=download[^"]+)', contents) | |
if m: | |
url = "https://docs.google.com" + m.groups()[0] | |
url = url.replace("&", "&") | |
return url | |
m = re.search(r'href="/open\?id=([^"]+)"', contents) | |
if m: | |
url = m.groups()[0] | |
uuid = re.search(r'<input\s+type="hidden"\s+name="uuid"\s+value="([^"]+)"', contents) | |
uuid = uuid.groups()[0] | |
url = "https://drive.usercontent.google.com/download?id=" + url + "&confirm=t&uuid=" + uuid | |
return url | |
m = re.search(r'"downloadUrl":"([^"]+)', contents) | |
if m: | |
url = m.groups()[0] | |
url = url.replace("\\u003d", "=") | |
url = url.replace("\\u0026", "&") | |
return url | |
m = re.search(r'<p class="uc-error-subcaption">(.*)</p>', contents) | |
if m: | |
error = m.groups()[0] | |
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): | |
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) | |
# Load cookies if exists | |
cookies_file = osp.join(home, ".cache/gdown/cookies.json") | |
if osp.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 | |
if return_cookies_file: | |
return sess, cookies_file | |
else: | |
return 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 = "https://drive.google.com/uc?id={id}".format(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 = "https://drive.google.com/uc?id={id}".format(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 = "https://drive.google.com/open?id={id}".format( | |
id=gdrive_file_id | |
) | |
continue | |
if res.headers["Content-Type"].startswith("text/html"): | |
m = re.search("<title>(.+)</title>", res.text) | |
if m and m.groups()[0].endswith(" - Google Docs"): | |
url = ( | |
"https://docs.google.com/document/d/{id}/export" | |
"?format={format}".format( | |
id=gdrive_file_id, | |
format="docx" if format is None else format, | |
) | |
) | |
continue | |
elif m and m.groups()[0].endswith(" - Google Sheets"): | |
url = ( | |
"https://docs.google.com/spreadsheets/d/{id}/export" | |
"?format={format}".format( | |
id=gdrive_file_id, | |
format="xlsx" if format is None else format, | |
) | |
) | |
continue | |
elif m and m.groups()[0].endswith(" - Google Slides"): | |
url = ( | |
"https://docs.google.com/presentation/d/{id}/export" | |
"?format={format}".format( | |
id=gdrive_file_id, | |
format="pptx" if format is None else format, | |
) | |
) | |
continue | |
elif ( | |
"Content-Disposition" in res.headers | |
and res.headers["Content-Disposition"].endswith("pptx") | |
and format not in {None, "pptx"} | |
): | |
url = ( | |
"https://docs.google.com/presentation/d/{id}/export" | |
"?format={format}".format( | |
id=gdrive_file_id, | |
format="pptx" if format is None else format, | |
) | |
) | |
continue | |
if use_cookies: | |
if not osp.exists(osp.dirname(cookies_file)): | |
os.makedirs(osp.dirname(cookies_file)) | |
# Save cookies | |
with open(cookies_file, "w") as f: | |
cookies = [ | |
(k, v) | |
for k, v in sess.cookies.items() | |
if not k.startswith("download_warning_") | |
] | |
json.dump(cookies, f, indent=2) | |
if "Content-Disposition" in res.headers: | |
# This is the file | |
break | |
if not (gdrive_file_id and is_gdrive_download_link): | |
break | |
# Need to redirect with confirmation | |
try: | |
url = get_url_from_gdrive_confirmation(res.text) | |
except FileURLRetrievalError as e: | |
message = ( | |
"Failed to retrieve file url:\n\n{}\n\n" | |
"You may still be able to access the file from the browser:" | |
"\n\n\t{}\n\n" | |
"but Gdown can't. Please check connections and permissions." | |
).format( | |
indent("\n".join(textwrap.wrap(str(e))), prefix="\t"), | |
url_origin, | |
) | |
raise FileURLRetrievalError(message) | |
if gdrive_file_id and is_gdrive_download_link: | |
content_disposition = six.moves.urllib_parse.unquote( | |
res.headers["Content-Disposition"] | |
) | |
m = re.search(r"filename\*=UTF-8''(.*)", content_disposition) | |
if not m: | |
m = re.search(r'filename=["\']?(.*?)["\']?$', content_disposition) | |
filename_from_url = m.groups()[0] | |
filename_from_url = filename_from_url.replace(osp.sep, "_") | |
else: | |
filename_from_url = osp.basename(url) | |
if output is None: | |
output = filename_from_url | |
output_is_path = isinstance(output, six.string_types) | |
if output_is_path and output.endswith(osp.sep): | |
if not osp.exists(output): | |
os.makedirs(output) | |
output = osp.join(output, filename_from_url) | |
if output_is_path: | |
existing_tmp_files = [] | |
for file in os.listdir(osp.dirname(output) or "."): | |
if file.startswith(osp.basename(output)): | |
existing_tmp_files.append(osp.join(osp.dirname(output), file)) | |
if resume and existing_tmp_files: | |
if len(existing_tmp_files) != 1: | |
print( | |
"There are multiple temporary files to resume:", | |
file=sys.stderr, | |
) | |
print("\n") | |
for file in existing_tmp_files: | |
print("\t", file, file=sys.stderr) | |
print("\n") | |
print( | |
"Please remove them except one to resume downloading.", | |
file=sys.stderr, | |
) | |
return | |
tmp_file = existing_tmp_files[0] | |
else: | |
resume = False | |
# mkstemp is preferred, but does not work on Windows | |
# https://github.com/wkentaro/gdown/issues/153 | |
tmp_file = tempfile.mktemp( | |
suffix=tempfile.template, | |
prefix=osp.basename(output), | |
dir=osp.dirname(output), | |
) | |
f = open(tmp_file, "ab") | |
else: | |
tmp_file = None | |
f = output | |
if tmp_file is not None and f.tell() != 0: | |
headers = {"Range": "bytes={}-".format(f.tell())} | |
res = sess.get(url, headers=headers, stream=True, verify=verify) | |
if not quiet: | |
# print("Downloading...", file=sys.stderr) | |
if resume: | |
print("Resume:", tmp_file, file=sys.stderr) | |
# if url_origin != url: | |
# print("From (original):", url_origin, file=sys.stderr) | |
# print("From (redirected):", url, file=sys.stderr) | |
# else: | |
# print("From:", url, file=sys.stderr) | |
print( | |
"To:", | |
osp.abspath(output) if output_is_path else output, | |
file=sys.stderr, | |
) | |
try: | |
total = res.headers.get("Content-Length") | |
if total is not None: | |
total = int(total) | |
if not quiet: | |
pbar = tqdm.tqdm(total=total, unit="B", unit_scale=True) | |
t_start = time.time() | |
for chunk in res.iter_content(chunk_size=CHUNK_SIZE): | |
f.write(chunk) | |
if not quiet: | |
pbar.update(len(chunk)) | |
if speed is not None: | |
elapsed_time_expected = 1.0 * pbar.n / speed | |
elapsed_time = time.time() - t_start | |
if elapsed_time < elapsed_time_expected: | |
time.sleep(elapsed_time_expected - elapsed_time) | |
if not quiet: | |
pbar.close() | |
if tmp_file: | |
f.close() | |
shutil.move(tmp_file, output) | |
finally: | |
sess.close() | |
return output | |