modify pytaxonkit.py to handle taxonkit's fuzzy matching

pytaxonkit.py

@@ -102,6 +102,17 @@ def validate_threads(value):
         return None
 
 
+def validate_n(value):
+    if value is None:
+        return None
+    try:
+        n = int(value)
+        return str(n)
+    except ValueError:
+        log(f'invalid n count "{value}"; resetting to taxonkit default', level="warning")
+        return None
+
+
 def test_validate_threads(capsys):
     assert validate_threads(None) is None
     assert validate_threads(2) == "2"
@@ -578,7 +589,9 @@ def test_name_empty():
 # -------------------------------------------------------------------------------------------------
 
 
-def name2taxid(names, sciname=False, threads=None, data_dir=None, debug=False):
+def name2taxid(
+    names, sciname=False, threads=None, data_dir=None, debug=False, fuzzy=False, fuzzy_top_n=None
+):
     """query taxid by taxon scientific name
 
     Parameters
@@ -593,6 +606,10 @@ def name2taxid(names, sciname=False, threads=None, data_dir=None, debug=False):
     data_dir : str, default None
         Specify the location of the NCBI taxonomy `.dmp` files; by default, taxonkit searches in
         `~/.taxonkit/`
+    fuzzy: bool, default False
+        By default, name matches need to be exact; when 'fuzzy=True' fuzzy
+    fuzzy_top_n : int
+        Override the default taxonkit setting for number of matches in fuzzy search
     debug : bool, default False
         Print debugging output, e.g., system calls to `taxonkit`
 
@@ -627,6 +644,10 @@ def name2taxid(names, sciname=False, threads=None, data_dir=None, debug=False):
         arglist.extend(("--threads", validate_threads(threads)))
     if data_dir:
         arglist.extend(("--data-dir", validate_data_dir(data_dir)))  # pragma: no cover
+    if fuzzy:
+        arglist.append("--fuzzy")
+    if fuzzy_top_n:
+        arglist.extend(("--fuzzy-top-n", validate_n(fuzzy_top_n)))
     if debug:
         log(*arglist)  # pragma: no cover
     proc = Popen(arglist, stdin=PIPE, stdout=PIPE, stderr=PIPE, universal_newlines=True)
@@ -666,6 +687,15 @@ def test_name2taxid_empty():
         assert result is None
 
 
+def test_name2taxid_fuzzy():
+    result = name2taxid(["Paramecium tetraurelia strain Stock d4-2"], fuzzy=True)
+    print(result.to_string())
+    assert len(result) == 1
+    row = result.iloc[0]
+    assert row.TaxID == 412030
+    assert row.Rank == "strain"
+
+
 # -------------------------------------------------------------------------------------------------
 # taxonkit filter
 # -------------------------------------------------------------------------------------------------