Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support to an user defined progress bar inside download function. #241

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
20 changes: 17 additions & 3 deletions gdown/download.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ def download(
id=None,
fuzzy=False,
resume=False,
bar=None,
):
"""Download file from URL.

Expand Down Expand Up @@ -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
-------
Expand All @@ -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(
Expand Down Expand Up @@ -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)
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can make bar as the same api as tqdm.tqdm? and rename it to pbar in the argument?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @wkentaro. Do you mean have the attributes update and close like tqdm or just a function with one parameter?

Copy link
Owner

@wkentaro wkentaro Feb 14, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @jcfaracco. Yeah, I don't mind dropping tqdm if there is an ultimate API to support pbar as function argument.
I wonder how the other libraries are doing this.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we analyze progress, we can see a similar approach to tqdm but using different function names. Again, I'm not proposing to drop tqdm, just the ability to use a user-defined pbar. For example, mine is a ipywidget that unfortunately does not work well with tqdm :-( I would love to use gdown with my widget. :-)

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the example of progress. To understand more deeply about your needs, can you give me a short snippet of how you want to use with ipywidget?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure @wkentaro

This is some simple example

import threading

from IPython import display as disp
from ipywidgets import HBox, FloatProgress, Label

class NotebookProgressBar(threading.Thread):
    MIN_CUR = -2
    MIN_TOTAL = -1

    def __init__(self):
        threading.Thread.__init__(self)

        self.bar = None
        self.percentage = None
        self.data = None

        self.__lock = threading.Lock()
        self.__current = self.MIN_CUR
        self.__total = self.MIN_TOTAL
        self.__error = False

    def show(self):
        self.bar = FloatProgress(value=0, min=0, max=100)
        self.percentage = Label(value='0 %')
        self.data = Label(value='')
        box = HBox((self.percentage, self.bar, self.data))
        disp.display(box)

    def set_current(self, current, total):
        with self.__lock:
            self.__current = current
            self.__total = total

    def set_error(self, error):
        self.__error = error

    def run(self):
        while (not self.__error and self.__current < self.__total):
            time.sleep(1)

            if self.__current != self.MIN_CUR and self.__total != self.MIN_TOTAL:
                progress = (self.__current / self.__total) * 100
                self.bar.value = progress
                self.percentage.value = f"{int(self.bar.value)} %%"
                self.data.value = f"{int(self.__current)} / {int(self.__total)}"

        if not self.__error:
            self.bar.style.bar_color = '#03c04a'
        else:
            self.bar.style.bar_color = '#ff0000'

All I need to do is update the bar with set_current() method. I usually use with wget, but this module does not work properly with Google Files for obvious reasons.

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