webmap-download: Refactor download().

1 files changed, 27 insertions, 29 deletions

diff --git a/webmap-download b/webmap-download
index 68f7b14..0165462 100755
--- a/webmap-download
+++ b/webmap-download
@@ -64,19 +64,22 @@ def download(url, dest, dir_fd=None, headers={}, session=requests, progress=None
             logging.exception('Could not parse Last-Modified value')
             last_modified = None
 
+    size = 0
+    pbar = None
+
     # XXX we can't use TemporaryFile as it uses O_EXCL, cf.
     # https://discuss.python.org/t/temporaryfile-contextmanager-that-allows-creating-a-directory-entry-on-success/19094/2
     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
-
+    try:
         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:
+            pbar = progress(
+                total=int(body_size) if body_size is not None else float('inf'),
+                leave=False,
+                unit_scale=True,
+                unit_divisor=1024,
+                unit='B'
+            )
+        with os.fdopen(fd, mode='wb', closefd=False) as fp:
             for chunk in r.iter_content(chunk_size=2**16):
                 chunk_size = len(chunk)
                 if pbar is not None:
@@ -85,36 +88,31 @@ def download(url, dest, dir_fd=None, headers={}, session=requests, progress=None
                 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
-        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)))
+
+        if last_modified is not None:
+            os.utime(fd, times=(last_modified, last_modified), follow_symlinks=True)
 
         # 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
         # (shorter race) than if we were dumping chunks in a named file descriptor
-        os.link(f'/proc/self/fd/{fp.fileno()}', dest_tmp,
-            dst_dir_fd=dir_fd, follow_symlinks=True)
-
-    # no need to close fd here, it was taken care of by the context manager above
+        os.link(f'/proc/self/fd/{fd}', dest_tmp, dst_dir_fd=dir_fd, follow_symlinks=True)
+    finally:
+        os.close(fd)
+        if pbar is not None:
+            pbar.close()
 
     try:
-        if last_modified is not None:
-            # XXX os.utime() doesn't work on file descriptors so we set mtime
-            # after linkat() instead
-            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)
-    except Exception as e:
+    except (OSError, ValueError) as e:
         try:
             os.unlink(dest_tmp, dir_fd=dir_fd)
-        except Exception:
-            pass
-        raise e
+        finally:
+            raise e
+
+    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)))
 
 def format_bytes(n):
     if n < 768: