diff --git a/gdown/download.py b/gdown/download.py index a62d2505..c65fe3cf 100644 --- a/gdown/download.py +++ b/gdown/download.py @@ -92,6 +92,7 @@ def download( id=None, fuzzy=False, resume=False, + bar=None, ): """Download file from URL. @@ -120,6 +121,9 @@ def download( resume: bool Resume the download from existing tmp file if possible. Default is False. + bar: type + A user defined progress bar. It should be followed by the arguments + current size and total size respectively. Returns ------- @@ -131,6 +135,9 @@ def download( if id is not None: url = "https://drive.google.com/uc?id={id}".format(id=id) + if bar is not None and not callable(bar): + raise ValueError("Parameter bar should be callable and respect arguments of cur_size and total_size") + url_origin = url sess, cookies_file = _get_session( @@ -265,19 +272,26 @@ def download( total = res.headers.get("Content-Length") if total is not None: total = int(total) - if not quiet: + if not quiet and bar is None: pbar = tqdm.tqdm(total=total, unit="B", unit_scale=True) + t_start = time.time() + + chunk_count = 0 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: + chunk_count += len(chunk) + bar(chunk_count, total) 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: + if not quiet and bar is None: pbar.close() if tmp_file: f.close()