From 48905f968f5d275186fef352e23935f3b9b6e3b7 Mon Sep 17 00:00:00 2001
From: Harrison Deng <yunyangdeng@outlook.com>
Date: Thu, 27 Feb 2025 04:55:43 +0000
Subject: [PATCH] Fixed scheme name argument not converting to ID

Added additional error messages to help user with usage.
---
 src/autobigs/cli/st.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/autobigs/cli/st.py b/src/autobigs/cli/st.py
index ec115a1..636f71c 100644
--- a/src/autobigs/cli/st.py
+++ b/src/autobigs/cli/st.py
@@ -57,17 +57,28 @@ async def run(args: Namespace):
     async with BIGSdbIndex() as bigsdb_index:
         gen_strings = read_multiple_fastas(args.fastas)
         scheme_id_lookup = await bigsdb_index.get_schemes_for_seqdefdb(args.seqdefdb)
-        scheme_id = args.scheme_id or args.scheme_name or (scheme_id_lookup)["MLST"]
         scheme_name_lookup = {value: key for key, value in scheme_id_lookup.items()}
+        known_dbs = await bigsdb_index.get_known_seqdef_dbs()
 
-        async with await bigsdb_index.build_profiler_from_seqdefdb(False, args.seqdefdb, scheme_id) as mlst_profiler:
+        if args.seqdefdb not in known_dbs:
+            raise ValueError("\"{0}\" is a known database. See -h for help.")
+        if args.scheme_id and args.scheme_id not in scheme_id_lookup.values():
+            raise ValueError("ID {0} not a known database scheme ID for database \"{1}\". See -h for help.".format(args.scheme_id, args.seqdefdb))
+        if args.scheme_name and args.scheme_name not in scheme_id_lookup:
+            raise ValueError("\"{0}\" not a known database scheme name for database \"{1}\". See -h for help.".format(args.scheme_name, args.seqdefdb))
+        if not (args.scheme_name or args.scheme_id) and "MLST" not in scheme_id_lookup:
+            raise ValueError("\"MLST\" not a known database scheme name for database \"{0}\". See -h for help.".format(args.seqdefdb))
+
+        selected_scheme_id = args.scheme_id or (scheme_id_lookup[args.scheme_name] if args.scheme_name else None) or scheme_id_lookup["MLST"]
+
+        async with await bigsdb_index.build_profiler_from_seqdefdb(False, args.seqdefdb, selected_scheme_id) as mlst_profiler:
             if not isinstance(mlst_profiler, BIGSdbMLSTProfiler):
                 raise TypeError("MLST profiler type invalid")
             mlst_profiles = mlst_profiler.profile_multiple_strings(gen_strings, args.stop_on_fail)
             failed = await write_mlst_profiles_as_csv(mlst_profiles, args.out)
             if len(failed) > 0:
                 print(f"A total of {len(failed)} IDs failed (no profile found):\n{"\n".join(failed)}")
-            print(f"Completed fetching from {args.seqdefdb} for {scheme_name_lookup[scheme_id]}s for {len(args.fastas)} sequences.")
+            print(f"Completed fetching from {args.seqdefdb} for {scheme_name_lookup[selected_scheme_id]}s for {len(args.fastas)} sequences.")
 
 def run_asynchronously(args):
     asyncio.run(run(args))