aboutsummaryrefslogtreecommitdiffstats
path: root/webmap-download
diff options
context:
space:
mode:
Diffstat (limited to 'webmap-download')
-rwxr-xr-xwebmap-download45
1 files changed, 33 insertions, 12 deletions
diff --git a/webmap-download b/webmap-download
index 39c48c2..9d0861a 100755
--- a/webmap-download
+++ b/webmap-download
@@ -30,7 +30,7 @@ def download_trystream(url, **kwargs):
r.raise_for_status()
return r
-def download(url, dest, dir_fd=None, headers={}, max_size=None, session=requests):
+def download(url, dest, dir_fd=None, headers={}, max_size=None, session=requests, progress=None):
logging.info('Downloading %s…', url)
destPath = Path(dest)
dest_tmp = destPath.with_stem(f'.{destPath.stem}.new').as_posix()
@@ -49,6 +49,7 @@ def download(url, dest, dir_fd=None, headers={}, max_size=None, session=requests
logging.info('%s: %d Not Modified', dest, r.status_code)
return
+ body_size = r.headers.get('Content-Length', None)
last_modified = r.headers.get('Last-Modified', None)
if last_modified is not None:
try:
@@ -63,13 +64,30 @@ def download(url, dest, dir_fd=None, headers={}, max_size=None, session=requests
fd = os.open(os.path.dirname(dest), O_WRONLY|O_CLOEXEC|O_TMPFILE, mode=0o644, dir_fd=dir_fd)
with os.fdopen(fd, mode='wb') as fp:
size = 0
- for chunk in r.iter_content(chunk_size=2**16):
- size = size + len(chunk)
- if max_size is not None and size > max_size:
- raise Exception(f'Payload exceeds max-size ({max_size})')
- fp.write(chunk)
+
+ if progress is not None:
+ tot = int(body_size) if body_size is not None else float('inf')
+ pbar = progress(total=tot, leave=False, unit_scale=True, unit_divisor=1024, unit='B')
+ else:
+ pbar = None
+
+ try:
+ for chunk in r.iter_content(chunk_size=2**16):
+ chunk_size = len(chunk)
+ if pbar is not None:
+ pbar.update(chunk_size)
+ size = size + chunk_size
+ if max_size is not None and size > max_size:
+ raise Exception(f'Payload exceeds max-size ({max_size})')
+ fp.write(chunk)
+ finally:
+ if pbar is not None:
+ pbar.close()
+
r = None
- end = time_monotonic()
+ elapsed = time_monotonic() - start
+ logging.info("%s: Downloaded %s in %s (%s/s)", dest, format_bytes(size),
+ format_time(elapsed), format_bytes(int(size/elapsed)))
# XXX unfortunately there is no way for linkat() to clobber the destination,
# so we use a temporary file; it's racy, but thanks to O_TMPFILE better
@@ -86,10 +104,6 @@ def download(url, dest, dir_fd=None, headers={}, max_size=None, session=requests
os.utime(dest_tmp, times=(last_modified, last_modified),
dir_fd=dir_fd, follow_symlinks=False)
os.rename(dest_tmp, dest, src_dir_fd=dir_fd, dst_dir_fd=dir_fd)
-
- elapsed = end - start
- logging.info("%s: Downloaded %s in %s (%s/s)", dest, format_bytes(size),
- format_time(elapsed), format_bytes(int(size/elapsed)))
except Exception as e:
try:
os.unlink(dest_tmp, dir_fd=dir_fd)
@@ -144,6 +158,12 @@ if __name__ == '__main__':
sourcedef['layername'] = name
sources.append(sourcedef)
+ if sys.stderr.isatty():
+ from tqdm import tqdm
+ pbar = tqdm
+ else:
+ pbar = None
+
# intentionally leave the dirfd open until the program terminates
opendir_args = O_RDONLY|O_CLOEXEC|O_PATH|O_DIRECTORY
destdir_fd = None if args.cachedir is None else os.open(args.cachedir, opendir_args)
@@ -224,7 +244,8 @@ if __name__ == '__main__':
headers['If-Modified-Since'] = formatdate(timeval=st.st_mtime, localtime=False, usegmt=True)
max_size = dl.get('max-size', 2**26) # 64MiB
download(url, dest, dir_fd=destdir_fd, max_size=max_size,
- headers=headers, session=sessionRequests)
+ headers=headers, session=sessionRequests,
+ progress=pbar)
downloads.add(url)
finally:
os.close(lock_fd)