webmap-download: Add a progress bar.

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)