Spaces:
Runtime error
Runtime error
File size: 3,051 Bytes
69a6cef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import os
from contextlib import contextmanager
import requests
from tqdm.auto import tqdm
from .session import get_requests_session, srequest
class _FakeClass:
def update(self, *args, **kwargs):
pass
@contextmanager
def _with_tqdm(expected_size, desc, silent: bool = False):
"""
Context manager that provides a tqdm progress bar for tracking the download progress.
:param expected_size: The expected size of the file being downloaded.
:type expected_size: int
:param desc: The description of the progress bar.
:type desc: str
:param silent: Whether to silence the progress bar. If True, a fake progress bar is used. (default: False)
:type silent: bool
"""
if not silent:
with tqdm(total=expected_size, unit='B', unit_scale=True, unit_divisor=1024, desc=desc) as pbar:
yield pbar
else:
yield _FakeClass()
def download_file(url, filename, expected_size: int = None, desc=None, session=None, silent: bool = False, **kwargs):
"""
Downloads a file from the given URL and saves it to the specified filename.
:param url: The URL of the file to download.
:type url: str
:param filename: The filename to save the downloaded file to.
:type filename: str
:param expected_size: The expected size of the file in bytes. (default: None)
:type expected_size: int
:param desc: The description of the download progress. If not provided, the filename is used. (default: None)
:type desc: str
:param session: An existing requests Session object to use for the download. If not provided, a new Session object is created. (default: None)
:type session: requests.Session
:param silent: Whether to silence the progress bar. If True, no progress bar is displayed. (default: False)
:type silent: bool
:param kwargs: Additional keyword arguments to pass to the `srequest` function.
:type kwargs: dict
:returns: The filename of the downloaded file.
:rtype: str
"""
session = session or get_requests_session()
response = srequest(session, 'GET', url, stream=True, allow_redirects=True, **kwargs)
expected_size = expected_size or response.headers.get('Content-Length', None)
expected_size = int(expected_size) if expected_size is not None else expected_size
desc = desc or os.path.basename(filename)
directory = os.path.dirname(filename)
if directory:
os.makedirs(directory, exist_ok=True)
with open(filename, 'wb') as f:
with _with_tqdm(expected_size, desc, silent) as pbar:
for chunk in response.iter_content(chunk_size=1024):
f.write(chunk)
pbar.update(len(chunk))
actual_size = os.path.getsize(filename)
if expected_size is not None and actual_size != expected_size:
os.remove(filename)
raise requests.exceptions.HTTPError(f"Downloaded file is not of expected size, "
f"{expected_size} expected but {actual_size} found.")
return filename
|