what.utils.file

 1import os.path
 2import hashlib
 3
 4import progressbar
 5import urllib.request
 6
 7class __MyProgressBar__():
 8    def __init__(self):
 9        self.pbar = None
10
11    def __call__(self, block_num, block_size, total_size):
12        if not self.pbar:
13            self.pbar=progressbar.ProgressBar(maxval=total_size)
14            self.pbar.start()
15
16        downloaded = block_num * block_size
17        if downloaded < total_size:
18            self.pbar.update(downloaded)
19        else:
20            self.pbar.finish()
21
22def download_file(file_name, path, url, hash):
23    print('Downloading', os.path.join(path, file_name))
24
25    if not os.path.exists(path):
26        os.makedirs(path)
27
28    try:
29        try:
30            urllib.request.urlretrieve(url, os.path.join(path, file_name), __MyProgressBar__())
31        except Exception as e:
32            print(e)
33    except (Exception, KeyboardInterrupt):
34        if os.path.exists(os.path.join(path, file_name)):
35            os.remove(os.path.join(path, file_name))
36        raise
37
38    # Validate download if succeeded and user provided an expected hash
39    sha256_hash = hashlib.sha256()
40    with open(os.path.join(path, file_name),"rb") as f:
41        for byte_block in iter(lambda: f.read(4096),b""):
42            sha256_hash.update(byte_block)
43        if (sha256_hash.hexdigest() == hash):
44            print("Model {} downloaded".format(os.path.join(path, file_name)))
45        else:
46            print('Incomplete or corrupted file detected.')
47
48def get_file(file_name, path, url, hash):
49    if os.path.isfile(os.path.join(path, file_name)):
50        sha256_hash = hashlib.sha256()
51        with open(os.path.join(path, file_name),"rb") as f:
52            for byte_block in iter(lambda: f.read(4096),b""):
53                sha256_hash.update(byte_block)
54            if (sha256_hash.hexdigest() == hash):
55                print("Model {} already exists".format(os.path.join(path, file_name)))
56            else:
57                redownload = input("File corrupted, Redownload? [Y/n]")
58                if len(redownload) == 0:
59                    redownload = 'y'
60                else:
61                    while not redownload in ['y', 'n']:
62                        redownload = input("File corrupted, Redownload? [y or n]")
63
64                if redownload == 'y':
65                    download_file(file_name, path, url, hash)
66
67    else:
68        download_file(file_name, path, url, hash)