Spaces:
Running
Running
import os | |
import re | |
import sys | |
import json | |
import time | |
import shutil | |
import tempfile | |
import warnings | |
from typing import Optional, Union, IO | |
import requests | |
from urllib.parse import urlparse, unquote | |
from tqdm import tqdm | |
CHUNK_SIZE = 512 * 1024 | |
HOME = os.path.expanduser("~") | |
def indent(text: str, prefix: str): | |
"""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): | |
"""Custom exception for issues retrieving file URLs.""" | |
def _extract_download_url_from_confirmation(contents: str, url_origin: str): | |
"""Extract the download URL from a Google Drive confirmation page.""" | |
patterns = [ | |
r'href="(\/uc\?export=download[^"]+)', | |
r'href="/open\?id=([^"]+)"', | |
r'"downloadUrl":"([^"]+)', | |
] | |
for pattern in patterns: | |
match = re.search(pattern, contents) | |
if match: | |
url = match.group(1) | |
if pattern == r'href="/open\?id=([^"]+)"': | |
uuid_match = re.search( | |
r'<input\s+type="hidden"\s+name="uuid"\s+value="([^"]+)"', | |
contents, | |
) | |
if uuid_match: | |
uuid = uuid_match.group(1) | |
return ( | |
"https://drive.usercontent.google.com/download?id=" | |
+ url | |
+ "&confirm=t&uuid=" | |
+ uuid | |
) | |
raise FileURLRetrievalError( | |
f"Could not find UUID for download from {url_origin}" | |
) | |
elif pattern == r'"downloadUrl":"([^"]+)': | |
return url.replace("\\u003d", "=").replace("\\u0026", "&") | |
else: | |
return "https://docs.google.com" + url.replace("&", "&") | |
error_match = re.search(r'<p class="uc-error-subcaption">(.*)</p>', contents) | |
if error_match: | |
error = 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 _create_session( | |
proxy: Optional[str] = None, | |
use_cookies: bool = True, | |
return_cookies_file: bool = 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: | |
sess.proxies = {"http": proxy, "https": proxy} | |
cookies_file = os.path.join(HOME, ".cache/gdown/cookies.json") | |
if os.path.exists(cookies_file) and use_cookies: | |
try: | |
with open(cookies_file) as f: | |
cookies = json.load(f) | |
for k, v in cookies: | |
sess.cookies[k] = v | |
except json.JSONDecodeError: | |
warnings.warn("Corrupted Cookies file") | |
return (sess, cookies_file) if return_cookies_file else sess | |
def download( | |
output: Optional[str] = None, | |
quiet: bool = False, | |
proxy: Optional[str] = None, | |
speed: Optional[float] = None, | |
use_cookies: bool = True, | |
verify: Union[bool, str] = True, | |
id: Optional[str] = None, | |
fuzzy: bool = True, | |
resume: bool = False, | |
format: Optional[str] = None, | |
url: Optional[str] = None, | |
): | |
"""Download a file from a URL, supporting Google Drive links. | |
Args: | |
output: Output filepath. Default is basename of URL. | |
quiet: Suppress terminal output. | |
proxy: HTTP/HTTPS proxy. | |
speed: Download speed limit (bytes per second). | |
use_cookies: Flag to use cookies. | |
verify: Verify TLS certificates. | |
id: Google Drive's file ID. | |
fuzzy: Fuzzy Google Drive ID extraction. | |
resume: Resume download from a tmp file. | |
format: Format for Google Docs/Sheets/Slides. | |
url: URL to download from. | |
Returns: | |
Output filename, or None on error. | |
""" | |
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 = _create_session( | |
proxy=proxy, use_cookies=use_cookies, return_cookies_file=True | |
) | |
while True: | |
res = sess.get(url, stream=True, verify=verify) | |
res.raise_for_status() | |
if url == url_origin and res.status_code == 500: | |
url = f"https://drive.google.com/open?id={id}" | |
continue | |
if res.headers.get("Content-Type", "").startswith("text/html"): | |
title_match = re.search("<title>(.+)</title>", res.text) | |
if title_match: | |
title = title_match.group(1) | |
if title.endswith(" - Google Docs"): | |
url = f"https://docs.google.com/document/d/{id}/export?format={'docx' if format is None else format}" | |
continue | |
if title.endswith(" - Google Sheets"): | |
url = f"https://docs.google.com/spreadsheets/d/{id}/export?format={'xlsx' if format is None else format}" | |
continue | |
if title.endswith(" - Google Slides"): | |
url = f"https://docs.google.com/presentation/d/{id}/export?format={'pptx' if format is None else format}" | |
continue | |
if ( | |
"Content-Disposition" in res.headers | |
and res.headers["Content-Disposition"].endswith("pptx") | |
and format not in (None, "pptx") | |
): | |
url = f"https://docs.google.com/presentation/d/{id}/export?format={'pptx' if format is None else format}" | |
continue | |
if use_cookies: | |
os.makedirs(os.path.dirname(cookies_file), exist_ok=True) | |
cookies = [ | |
(k, v) | |
for k, v in sess.cookies.items() | |
if not k.startswith("download_warning_") | |
] | |
with open(cookies_file, "w") as f: | |
json.dump(cookies, f, indent=2) | |
if "Content-Disposition" in res.headers: | |
break | |
parsed_url = urlparse(url) | |
is_gdrive = parsed_url.hostname in ("drive.google.com", "docs.google.com") | |
is_download_link = parsed_url.path.endswith("/uc") | |
if not (is_gdrive and is_download_link and fuzzy): | |
break | |
try: | |
url = _extract_download_url_from_confirmation(res.text, url_origin) | |
except FileURLRetrievalError as e: | |
raise FileURLRetrievalError(e) | |
content_disposition = res.headers.get("Content-Disposition", "") | |
filename_match = re.search( | |
r"filename\*=UTF-8''(.*)", content_disposition | |
) or re.search(r'filename=["\']?(.*?)["\']?$', content_disposition) | |
filename_from_url = ( | |
unquote(filename_match.group(1)) if filename_match else os.path.basename(url) | |
) | |
download_path = output or filename_from_url | |
if isinstance(download_path, str) and download_path.endswith(os.path.sep): | |
os.makedirs(download_path, exist_ok=True) | |
download_path = os.path.join(download_path, filename_from_url) | |
temp_dir = os.path.dirname(download_path) or "." | |
prefix = os.path.basename(download_path) | |
if isinstance(download_path, str): | |
existing_tmp_files = [ | |
os.path.join(temp_dir, file) | |
for file in os.listdir(temp_dir) | |
if file.startswith(prefix) | |
] | |
if resume and existing_tmp_files: | |
if len(existing_tmp_files) > 1: | |
print( | |
"There are multiple temporary files to resume:", | |
file=sys.stderr, | |
) | |
for file in existing_tmp_files: | |
print(f"\t{file}", file=sys.stderr) | |
print( | |
"Please remove them except one to resume downloading.", | |
file=sys.stderr, | |
) | |
return None | |
temp_file_path = existing_tmp_files[0] | |
else: | |
resume = False | |
temp_file_path = tempfile.mktemp( | |
suffix=tempfile.template, prefix=prefix, dir=temp_dir | |
) | |
try: | |
file_obj: IO = open(temp_file_path, "ab") | |
except Exception as e: | |
print( | |
f"Could not open the temporary file {temp_file_path}: {e}", | |
file=sys.stderr, | |
) | |
return None | |
else: | |
temp_file_path = None | |
file_obj = download_path | |
if temp_file_path is not None and file_obj.tell() != 0: | |
headers = {"Range": f"bytes={file_obj.tell()}-"} | |
res = sess.get(url, headers=headers, stream=True, verify=verify) | |
res.raise_for_status() | |
try: | |
total = int(res.headers.get("Content-Length", 0)) | |
if total > 0: | |
if not quiet: | |
pbar = tqdm( | |
total=total, unit="B", unit_scale=True, desc=filename_from_url | |
) | |
else: | |
if not quiet: | |
pbar = tqdm(unit="B", unit_scale=True, desc=filename_from_url) | |
t_start = time.time() | |
for chunk in res.iter_content(chunk_size=CHUNK_SIZE): | |
file_obj.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 temp_file_path: | |
file_obj.close() | |
shutil.move(temp_file_path, download_path) | |
finally: | |
sess.close() | |
return download_path | |