kyle0x54/medvision

import os
import shutil
from typing import Callable, Optional
from urllib import request


def download_url(
    url: str,
    dst_dir: str,
    *,
    filename: Optional[str] = None,
    progress: bool = True
) -> str:
    """
    Download a file from a given URL to a directory. If file exists, will not
        overwrite the existing file.

    Args:
        url (str):
        dst_dir (str): the directory to download the file
        filename (str or None): the basename to save the file.
            Will use the name in the URL if not given.
        progress (bool): whether to use tqdm to draw a progress bar.

    Returns:
        str: the path to the downloaded file or the existing one.
    """
    os.makedirs(dst_dir, exist_ok=True)
    if filename is None:
        filename = url.split("/")[-1]
        assert len(filename), "Cannot obtain filename from url {}".format(url)
    fpath = os.path.join(dst_dir, filename)

    if os.path.isfile(fpath):
        return fpath

    tmp = fpath + ".tmp"  # download to a tmp file first, to be more atomic.
    try:
        if progress:
            from tqdm import tqdm

            def hook(t: tqdm) -> Callable[[int, int, Optional[int]], None]:
                last_b = [0]

                def inner(
                    b: int, bsize: int, tsize: Optional[int] = None
                ) -> None:
                    if tsize is not None:
                        t.total = tsize
                    t.update((b - last_b[0]) * bsize)  # type: ignore
                    last_b[0] = b

                return inner

            with tqdm(unit="B", unit_scale=True, miniters=1, desc=filename,
                      leave=True) as t:
                tmp, _ = request.urlretrieve(
                    url, filename=tmp, reporthook=hook(t)
                )

        else:
            tmp, _ = request.urlretrieve(url, filename=tmp)
        statinfo = os.stat(tmp)
        size = statinfo.st_size
        if size == 0:
            raise IOError("Downloaded an empty file from {}!".format(url))
        # download to tmp first and move to fpath, to make this function more
        # atomic.
        shutil.move(tmp, fpath)
    except IOError:
        raise
    finally:
        try:
            os.unlink(tmp)
        except IOError:
            pass

    return fpath