From 34d3b890a361ca388a78361c58b578fc15b24852 Mon Sep 17 00:00:00 2001
From: Dan LaManna <dan.lamanna@kitware.com>
Date: Fri, 11 Feb 2022 12:17:50 -0500
Subject: [PATCH] Add collection filtering to download commands

---
 isic_cli/image.py    |  8 ++++++--
 isic_cli/metadata.py |  8 ++++++--
 isic_cli/utils.py    | 13 +++++++++----
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/isic_cli/image.py b/isic_cli/image.py
index 02f3e86..0926826 100644
--- a/isic_cli/image.py
+++ b/isic_cli/image.py
@@ -37,6 +37,10 @@ def _download_image(image: dict, to: Path, progress, task) -> None:
 def download_images(
     ctx: typer.Context,
     search: str = typer.Option(''),
+    collections: str = typer.Option(
+        '',
+        help='Limit the images based on a comma separated string of collection ids (see isic collection list).',  # noqa: E501
+    ),
     max_images: int = typer.Option(1_000, min=0, help='Use a value of 0 to disable the limit.'),
     outdir: Path = typer.Option(Path('images'), file_okay=False, dir_okay=True, writable=True),
 ):
@@ -56,12 +60,12 @@ def download_images(
     outdir.mkdir(exist_ok=True)
     with Progress() as progress:
         with get_session(ctx.obj.auth_headers) as session:
-            num_images = get_num_images(session, search)
+            num_images = get_num_images(session, search, collections)
             if max_images > 0:
                 num_images = min(num_images, max_images)
 
             task = progress.add_task(f'Downloading images ({num_images} total)', total=num_images)
-            images = get_images(session, search)
+            images = get_images(session, search, collections)
 
             if max_images > 0:
                 images = itertools.islice(images, max_images)
diff --git a/isic_cli/metadata.py b/isic_cli/metadata.py
index 56157d7..294f878 100644
--- a/isic_cli/metadata.py
+++ b/isic_cli/metadata.py
@@ -76,6 +76,10 @@ def validate_metadata(csv_file: Path):
 def download(
     ctx: typer.Context,
     search: str = typer.Option(''),
+    collections: str = typer.Option(
+        '',
+        help='Limit the images based on a comma separated string of collection ids (see isic collection list).',  # noqa: E501
+    ),
     max_results: int = typer.Option(1_000, min=0, help='Use a value of 0 to disable the limit.'),
 ):
     """
@@ -92,11 +96,11 @@ def download(
     anatom_site_general:*torso AND image_type:dermoscopic
     """
     with get_session(ctx.obj.auth_headers) as session:
-        num_results = get_num_images(session, search)
+        num_results = get_num_images(session, search, collections)
         if max_results > 0:
             num_results = min(num_results, max_results)
 
-        images = get_images(session, search)
+        images = get_images(session, search, collections)
 
         if max_results > 0:
             images = itertools.islice(images, max_results)
diff --git a/isic_cli/utils.py b/isic_cli/utils.py
index 9e12504..5e5e02a 100644
--- a/isic_cli/utils.py
+++ b/isic_cli/utils.py
@@ -13,8 +13,8 @@ def get_collections(session: IsicCliSession) -> Iterable[dict]:
         next_page = r.json()['next']
 
 
-def get_images(session: IsicCliSession, search: str) -> Iterable[dict]:
-    next_page = f'images/search/?query={search}'
+def get_images(session: IsicCliSession, search: str, collections: str) -> Iterable[dict]:
+    next_page = f'images/search/?query={search}' + f'&collections={collections}'
 
     while next_page:
         r = session.get(next_page)
@@ -23,7 +23,12 @@ def get_images(session: IsicCliSession, search: str) -> Iterable[dict]:
         next_page = r.json()['next']
 
 
-def get_num_images(session: IsicCliSession, search: str) -> int:
-    r = session.get(f'images/search/?query={search}&limit=1')
+def get_num_images(session: IsicCliSession, search: str, collections: str) -> int:
+    params = {
+        'query': search,
+        'collections': collections,
+        'limit': 1,
+    }
+    r = session.get('images/search/', params=params)
     r.raise_for_status()
     return r.json()['count']