download urls to tempfile then copy to cache
@@ -2,6 +2,8 @@ from typing import Tuple
import os
import base64
import logging
+import shutil
+import tempfile
from urllib.parse import urlparse
import requests
@@ -80,20 +82,29 @@ def get_from_cache(url: str, cache_dir: str = DATASET_CACHE) -> str:
etag = response.headers.get("ETag")
filename = url_to_filename(url, etag)
- path = os.path.join(cache_dir, filename)
+ # get cache path to put the file
+ cache_path = os.path.join(cache_dir, filename)
- if not os.path.exists(path):
- logger.info("%s not found in cache, downloading to %s", url, path)
+ if not os.path.exists(cache_path):
+ # Download to temporary file, then copy to cache dir once finished.
+ # Otherwise you get corrupt cache entries if the download gets interrupted.
+ _, temp_filename = tempfile.mkstemp()
+ logger.info("%s not found in cache, downloading to %s", url, temp_filename)
# GET file object
req = requests.get(url, stream=True)
content_length = req.headers.get('Content-Length')
total = int(content_length) if content_length is not None else None
progress = tqdm.tqdm(unit="B", total=total)
- with open(path, 'wb') as output_file:
+ with open(temp_filename, 'wb') as temp_file:
for chunk in req.iter_content(chunk_size=1024):
if chunk: # filter out keep-alive new chunks
progress.update(len(chunk))
- output_file.write(chunk)
+ temp_file.write(chunk)
- return path
+ logger.info("copying %s to cache at %s", temp_filename, cache_path)
+ shutil.copyfile(temp_filename, cache_path)
+ logger.info("removing temp file %s", temp_filename)
+ os.remove(temp_filename)
+
+ return cache_path
