| | import os |
| | import os.path |
| | import hashlib |
| | import errno |
| | from tqdm import tqdm |
| |
|
| |
|
| | def gen_bar_updater(pbar): |
| | def bar_update(count, block_size, total_size): |
| | if pbar.total is None and total_size: |
| | pbar.total = total_size |
| | progress_bytes = count * block_size |
| | pbar.update(progress_bytes - pbar.n) |
| |
|
| | return bar_update |
| |
|
| |
|
| | def check_integrity(fpath, md5=None): |
| | if md5 is None: |
| | return True |
| | if not os.path.isfile(fpath): |
| | return False |
| | md5o = hashlib.md5() |
| | with open(fpath, 'rb') as f: |
| | |
| | for chunk in iter(lambda: f.read(1024 * 1024), b''): |
| | md5o.update(chunk) |
| | md5c = md5o.hexdigest() |
| | if md5c != md5: |
| | return False |
| | return True |
| |
|
| |
|
| | def makedir_exist_ok(dirpath): |
| | """ |
| | Python2 support for os.makedirs(.., exist_ok=True) |
| | """ |
| | try: |
| | os.makedirs(dirpath) |
| | except OSError as e: |
| | if e.errno == errno.EEXIST: |
| | pass |
| | else: |
| | raise |
| |
|
| |
|
| | def download_url(url, root, filename=None, md5=None): |
| | """Download a file from a url and place it in root. |
| | Args: |
| | url (str): URL to download file from |
| | root (str): Directory to place downloaded file in |
| | filename (str): Name to save the file under. If None, use the basename of the URL |
| | md5 (str): MD5 checksum of the download. If None, do not check |
| | """ |
| | from six.moves import urllib |
| |
|
| | root = os.path.expanduser(root) |
| | if not filename: |
| | filename = os.path.basename(url) |
| | fpath = os.path.join(root, filename) |
| |
|
| | makedir_exist_ok(root) |
| |
|
| | |
| | if os.path.isfile(fpath) and check_integrity(fpath, md5): |
| | print('Using downloaded and verified file: ' + fpath) |
| | else: |
| | try: |
| | print('Downloading ' + url + ' to ' + fpath) |
| | urllib.request.urlretrieve( |
| | url, fpath, |
| | reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True)) |
| | ) |
| | except OSError: |
| | if url[:5] == 'https': |
| | url = url.replace('https:', 'http:') |
| | print('Failed download. Trying https -> http instead.' |
| | ' Downloading ' + url + ' to ' + fpath) |
| | urllib.request.urlretrieve( |
| | url, fpath, |
| | reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True)) |
| | ) |
| |
|
| |
|
| | def list_dir(root, prefix=False): |
| | """List all directories at a given root |
| | Args: |
| | root (str): Path to directory whose folders need to be listed |
| | prefix (bool, optional): If true, prepends the path to each result, otherwise |
| | only returns the name of the directories found |
| | """ |
| | root = os.path.expanduser(root) |
| | directories = list( |
| | filter( |
| | lambda p: os.path.isdir(os.path.join(root, p)), |
| | os.listdir(root) |
| | ) |
| | ) |
| |
|
| | if prefix is True: |
| | directories = [os.path.join(root, d) for d in directories] |
| |
|
| | return directories |
| |
|
| |
|
| | def list_files(root, suffix, prefix=False): |
| | """List all files ending with a suffix at a given root |
| | Args: |
| | root (str): Path to directory whose folders need to be listed |
| | suffix (str or tuple): Suffix of the files to match, e.g. '.png' or ('.jpg', '.png'). |
| | It uses the Python "str.endswith" method and is passed directly |
| | prefix (bool, optional): If true, prepends the path to each result, otherwise |
| | only returns the name of the files found |
| | """ |
| | root = os.path.expanduser(root) |
| | files = list( |
| | filter( |
| | lambda p: os.path.isfile(os.path.join(root, p)) and p.endswith(suffix), |
| | os.listdir(root) |
| | ) |
| | ) |
| |
|
| | if prefix is True: |
| | files = [os.path.join(root, d) for d in files] |
| |
|
| | return files |
