From a3ff68be5a911537038aec5f4b28cc80b876afe3 Mon Sep 17 00:00:00 2001 From: Dan Streetman Date: Thu, 16 Sep 2021 19:24:00 -0400 Subject: [PATCH] misc: download to tmp file, to avoid leftover 0-size file on error --- ubuntutools/misc.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/ubuntutools/misc.py b/ubuntutools/misc.py index 7719966..a8dc608 100644 --- a/ubuntutools/misc.py +++ b/ubuntutools/misc.py @@ -327,14 +327,17 @@ def download(src, dst, size=0): (src, username, password) = extract_authentication(src) auth = (username, password) if username or password else None - try: - with requests.get(src, stream=True, auth=auth) as fsrc, dst.open('wb') as fdst: - fsrc.raise_for_status() - _download(fsrc, fdst, size) - except requests.RequestException as e: - if e.response and e.response.status_code == 404: - raise NotFoundError(f'URL {src} not found') from e - raise DownloadError(f'Could not download {src} to {dst}') from e + with tempfile.TemporaryDirectory() as d: + tmpdst = Path(d) / 'dst' + try: + with requests.get(src, stream=True, auth=auth) as fsrc, tmpdst.open('wb') as fdst: + fsrc.raise_for_status() + _download(fsrc, fdst, size) + except requests.RequestException as e: + if e.response and e.response.status_code == 404: + raise NotFoundError(f'URL {src} not found') from e + raise DownloadError(f'Could not download {src} to {dst}') from e + shutil.move(tmpdst, dst) def _download(fsrc, fdst, size):