From 58120b4d8aaf17e9d678a5720341d1eb1217233d Mon Sep 17 00:00:00 2001
From: Julio Faracco <jcfaracco@gmail.com>
Date: Sat, 11 Feb 2023 15:04:05 -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 file and the total size.

Signed-off-by: Julio Faracco <jcfaracco@gmail.com>
---
 gdown/download.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

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()