Merge pull request #54 from ImageMarkup/threadpool-teardown

ImageMarkup · Jul 13, 2023 · 7af8138 · 7af8138
2 parents 5551301 + 2344a15
commit 7af8138
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/isic_cli/cli/image.py b/isic_cli/cli/image.py
@@ -104,17 +104,17 @@ def download(
             total=download_num_images,
         )
         # the futures ThreadPoolExecutor doesn't allow one to easily Ctrl-c
-        thread_pool = ThreadPool(max(10, os.cpu_count() or 10))
         images_iterator = itertools.islice(
             get_images(ctx.session, search, collections), download_num_images
         )
 
         # See comment above _extract_metadata for why this is necessary
         images = []
         func = functools.partial(download_image, to=outdir, progress=progress, task=task)
-        for image_chunk in chunked(images_iterator, 100):
-            images.extend(image_chunk)
-            thread_pool.map(func, image_chunk)
+        with ThreadPool(max(10, os.cpu_count() or 10)) as thread_pool:
+            for image_chunk in chunked(images_iterator, 100):
+                images.extend(image_chunk)
+                thread_pool.map(func, image_chunk)
 
         headers, records = _extract_metadata(images)
         with (outdir / "metadata.csv").open("w", encoding="utf8") as outfile: