diff options
| -rwxr-xr-x | webmap-download | 56 |
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: |