Tqdm Download Example

tqdm_download.py

"""
Tqdm Download Hook

Lei Mao
University of Chicago

Modified from https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py
"""

import os
from urllib.request import urlretrieve
from tqdm import tqdm

def my_hook(t):
    """Wraps tqdm instance.

    Don't forget to close() or __exit__()
    the tqdm instance once you're done with it (easiest using `with` syntax).

    Example
    -------

    >>> with tqdm(...) as t:
    ...     reporthook = my_hook(t)
    ...     urllib.urlretrieve(..., reporthook=reporthook)

    """
    last_b = [0]

    def update_to(b=1, bsize=1, tsize=None):
        """
        b  : int, optional
            Number of blocks transferred so far [default: 1].
        bsize  : int, optional
            Size of each block (in tqdm units) [default: 1].
        tsize  : int, optional
            Total size (in tqdm units). If [default: None] remains unchanged.
        """
        if tsize is not None:
            t.total = tsize
        t.update((b - last_b[0]) * bsize)
        last_b[0] = b

    return update_to


class TqdmUpTo(tqdm):
    """Alternative Class-based version of the above.

    Provides `update_to(n)` which uses `tqdm.update(delta_n)`.

    Inspired by [twine#242](https://github.com/pypa/twine/pull/242),
    [here](https://github.com/pypa/twine/commit/42e55e06).
    """

    def update_to(self, b=1, bsize=1, tsize=None):
        """
        b  : int, optional
            Number of blocks transferred so far [default: 1].
        bsize  : int, optional
            Size of each block (in tqdm units) [default: 1].
        tsize  : int, optional
            Total size (in tqdm units). If [default: None] remains unchanged.
        """
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)  # will also set self.n = b * bsize


def download1(url, save_dir):
    filename = url.split('/')[-1]
    with TqdmUpTo(unit = 'B', unit_scale = True, unit_divisor = 1024, miniters = 1, desc = filename) as t:
        urlretrieve(url, filename = os.path.join(save_dir, filename), reporthook = t.update_to)

def download2(url, save_dir):
    filename = url.split('/')[-1]
    with tqdm(unit = 'B', unit_scale = True, unit_divisor = 1024, miniters = 1, desc = filename) as t:
        urlretrieve(url, filename = os.path.join(save_dir, filename), reporthook = my_hook(t), data = None)


if __name__ == '__main__':

    url = 'https://caspersci.uk.to/matryoshka.zip'
    save_dir = './'

    # download1 and download2 are the same
    download1(url = url, save_dir = save_dir)
    download2(url = url, save_dir = save_dir)