Merge pull request #7 from ImageMarkup/clean-interrupt

Cleanup on keyboard interrupts

Author: danlamanna

isic_cli/cli/image.py

@@ -1,3 +1,4 @@
+import atexit
 import csv
 import itertools
 from pathlib import Path
@@ -16,6 +17,11 @@
 from isic_cli.io.http import download_image, get_images, get_num_images
 
 
+def cleanup_partially_downloaded_files(directory: Path) -> None:
+    for p in directory.glob('**/.isic-partial.*'):
+        p.unlink()
+
+
 @click.group(short_help='Manage images.')
 @click.pass_obj
 def image(ctx):
@@ -73,6 +79,8 @@ def download(
     """
     outdir.mkdir(exist_ok=True)
 
+    atexit.register(cleanup_partially_downloaded_files, outdir)
+
     with Progress(console=Console(file=sys.stderr)) as progress:
         archive_num_images = get_num_images(ctx.session, search, collections)
         download_num_images = archive_num_images if limit == 0 else min(archive_num_images, limit)

isic_cli/io/http.py

@@ -1,5 +1,7 @@
 import logging
 from pathlib import Path
+import shutil
+from tempfile import NamedTemporaryFile
 from typing import Iterable, Optional, Union
 
 from more_itertools import chunked
@@ -127,8 +129,12 @@ def download_image(image: dict, to: Path, progress, task) -> None:
         r = session.get(image['urls']['full'], stream=True)
         r.raise_for_status()
 
-        with open(to / f'{image["isic_id"]}.JPG', 'wb') as outfile:
+        temp_file_name = None
+        with NamedTemporaryFile(dir=to, prefix='.isic-partial.', delete=False) as outfile:
+            temp_file_name = outfile.name
             for chunk in r.iter_content(1024 * 1024 * 5):
                 outfile.write(chunk)
 
+        shutil.move(temp_file_name, to / f'{image["isic_id"]}.JPG')
+
     progress.update(task, advance=1)