MRG: set n-simultaneous-downloads to 9 if api key provided (#194)

Fixes #185
sourmash-bio · Jan 28, 2025 · 64d4132 · 64d4132
1 parent 5640003
commit 64d4132
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/python/sourmash_plugin_directsketch/__init__.py b/src/python/sourmash_plugin_directsketch/__init__.py
@@ -119,7 +119,7 @@ def __init__(self, p):
         p.add_argument(
             "-n",
             "--n-simultaneous-downloads",
-            default=3,
+            default=None,
             type=int,
             choices=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
             help="Number of accessions to download simultaneously (default=3). Must be <=3 if not using API key.",
@@ -162,7 +162,8 @@ def main(self, args):
             api_key = os.environ.get("NCBI_API_KEY", None)
             if api_key:
                 args.api_key = api_key
-            elif args.n_simultaneous_downloads > 3:
+            elif args.n_simultaneous_downloads is not None \
+                 and args.n_simultaneous_downloads > 3:
                 notify(
                     "Error: please provide an API Key to use n_simultaneous_downloads > 3."
                 )
@@ -176,6 +177,13 @@ def main(self, args):
 
         num_threads = set_thread_pool(args.cores)
 
+        if args.n_simultaneous_downloads is None:
+            if args.api_key:
+                notify("API key provided - setting --n-simultaneous-downloads to 9")
+                args.n_simultaneous_downloads = 9
+            else:
+                args.n_simultaneous_downloads = 3
+
         notify(
             f"Downloading and sketching all accessions in '{args.input_csv} using {args.n_simultaneous_downloads} simultaneous downloads, {args.retry_times} retries, and {num_threads} threads."
         )