Downloads a file and save it. Displays progressbar. Throws exception if sha256 hash does not match.

download_file.py

from hashlib import sha256
from tqdm.auto import tqdm
import logging

__author__ = "Jules Kreuer"
__copyright__ = "Jules Kreuer"
__license__ = "LGPL-3.0-only"

_logger = logging.getLogger(__name__)


def download_file(file_name, expected_hash):
    """
    Downloads a file and save it.
    Displays progressbar. Throws exception if sha256 hash does not match.
    Parameters
    ----------
        file_name: str, name of model file.
        expected_hash: str, sha256 hash of model.
    Returns:
    ----------
        True, Error if sha256 hash of model is not equal to hash.
    """
    _logger.info(f"Downloading: {file_name}")
    
    # TODO add url to model archive.
    base_url = "https://example.com/"
    file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "static", file_name)
    url = base_url + file_name
    
    r = requests.get(url, stream=True)
    total_size = int(r.headers.get("content-length", 0))
    with open(file_path, "wb") as file, tqdm(
        total=total_size, unit="B", unit_scale=True, unit_divisor=1024
    ) as pbar:
        for data in r.iter_content(1024):
            pbar.update(len(data))
            file.write(data)

    _logger.debug("Download complete.")
    _logger.debug("Checking now sha2567 hash.")

    # Check that file was downloaded properly.
    file_hash = sha256()
    with open(file_path, "rb") as f:
        # Read by block
        for block in iter(lambda: f.read(4096), b""):
            file_hash.update(block)
    if not file_hash.hexdigest() == expected_hash:
        raise AssertionError(
            f"File {file_name} was not downloaded properly. "
            "Use `TODO add` to delete the corrupt model files."
        )
    _logger.debug("sha2567 hash matches.")
    return True