From c08058e0f62dec0503bd0a5230e40a3ae25ad935 Mon Sep 17 00:00:00 2001 From: Julio Faracco Date: Sun, 29 Sep 2024 15:47:56 -0300 Subject: [PATCH] download: Add support to an user defined progress bar This commit just introduces a new argument to the download function: `bar`. The argument let the person who is developing code define his own progress bar using whatever he wants to. It also requires the pattern of passing the current size of the downloaded chunk. This callable function respects the tqdm API. Signed-off-by: Julio Faracco --- gdown/download.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/gdown/download.py b/gdown/download.py index baadef1..72eeff2 100644 --- a/gdown/download.py +++ b/gdown/download.py @@ -124,6 +124,7 @@ def download( format=None, user_agent=None, log_messages=None, + bar=None, ): """Download file from URL. @@ -165,6 +166,10 @@ def download( Log messages to customize. Currently it supports: - 'start': the message to show the start of the download - 'output': the message to show the output filename + bar: type + A user defined progress bar. It should be followed by the argument + of the current chunk size downloaded. The total size should be defined + previously in order to make the API consistent to tqdm. Returns ------- @@ -180,6 +185,8 @@ def download( user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36" # NOQA: E501 if log_messages is None: log_messages = {} + if bar is not None and not callable(bar): + raise ValueError("Parameter bar should be callable and respect arguments of cur_size") url_origin = url @@ -362,19 +369,25 @@ def download( total = res.headers.get("Content-Length") if total is not None: total = int(total) + start_size - if not quiet: + if not quiet and bar is None: pbar = tqdm.tqdm(total=total, unit="B", initial=start_size, unit_scale=True) + + percent = 0.0 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 bar is None: + pbar.update(len(chunk)) + else: + bar(len(chunk)) + percent += len(chunk) / total if speed is not None: - elapsed_time_expected = 1.0 * pbar.n / speed + elapsed_time_expected = 1.0 * percent / speed elapsed_time = time.time() - t_start if elapsed_time < elapsed_time_expected: time.sleep(elapsed_time_expected - elapsed_time) - if not quiet: + if not quiet and bar is None: pbar.close() if tmp_file: f.close()