Handle NCBI taxdump members with a leading './' (fixes #810)

ete4/ncbi_taxonomy/ncbiquery.py

@@ -585,6 +585,19 @@ def get_broken_branches(self, t, taxa_lineages, n2content=None):
         return broken_branches, broken_clades, broken_clade_sizes
 
 
+def extractfile_dmp(tar, name):
+    """Return the tar member `name`, tolerating a leading './' in its name.
+
+    Tarballs made with GNU tar (e.g. `tar -c -C dir .`) name their members
+    like './names.dmp', so a lookup for the bare 'names.dmp' fails. Try both.
+    See https://github.com/etetoolkit/ete/issues/810
+    """
+    try:
+        return tar.extractfile(name)
+    except KeyError:
+        return tar.extractfile('./' + name)
+
+
 def load_ncbi_tree_from_dump(tar):
     from .. import Tree
     parent2child = {}
@@ -595,7 +608,7 @@ def load_ncbi_tree_from_dump(tar):
     node2common = {}
     print("Loading node names...")
     unique_nocase_synonyms = set()
-    for line in tar.extractfile("names.dmp"):
+    for line in extractfile_dmp(tar, "names.dmp"):
         line = str(line.decode())
         fields =  [_f.strip() for _f in line.split("|")]
         nodename = fields[0]
@@ -622,7 +635,7 @@ def load_ncbi_tree_from_dump(tar):
     print(len(synonyms), "synonyms loaded.")
 
     print("Loading nodes...")
-    for line in tar.extractfile("nodes.dmp"):
+    for line in extractfile_dmp(tar, "nodes.dmp"):
         line = str(line.decode())
         fields =  line.split("|")
         nodename = fields[0].strip()
@@ -689,7 +702,7 @@ def update_db(dbfile, targz_file=None):
         SYN.write('\n'.join(["%s\t%s" %(v[0],v[1]) for v in synonyms]))
 
     with open("merged.tab", "w") as merged:
-        for line in tar.extractfile("merged.dmp"):
+        for line in extractfile_dmp(tar, "merged.dmp"):
             line = str(line.decode())
             out_line = '\t'.join([_f.strip() for _f in line.split('|')[:2]])
             merged.write(out_line+'\n')

tests/test_ncbiquery.py

@@ -2,7 +2,10 @@
 Test the functionality of ncbiquery.py. To run with pytest.
 """
 
+import io
 import os
+import tarfile
+
 import pytest
 
 from ete4 import PhyloTree, NCBITaxa, ETE_DATA_HOME, update_ete_data
@@ -166,6 +169,43 @@ def test_merged_id():
     assert t2 == [1, 131567, 2, 1783272, 1239, 186801, 3085636, 186803, 207244, 649756]
 
 
+def _make_taxdump_tar(path, prefix=''):
+    """Write a minimal taxdump tarball, prefixing member names with `prefix`.
+
+    GNU tar creates members named like `./names.dmp` when archiving a
+    directory with `tar -c -C dir .`, so we emulate that with prefix='./'.
+    """
+    names = ('1\t|\troot\t|\t\t|\tscientific name\t|\n'
+             '2\t|\tBacteria\t|\t\t|\tscientific name\t|\n')
+    nodes = ('1\t|\t1\t|\tno rank\t|\n'
+             '2\t|\t1\t|\tsuperkingdom\t|\n')
+    merged = '3\t|\t2\t|\n'
+
+    with tarfile.open(path, 'w:gz') as tar:
+        for fname, content in [('names.dmp', names),
+                               ('nodes.dmp', nodes),
+                               ('merged.dmp', merged)]:
+            data = content.encode()
+            info = tarfile.TarInfo(name=prefix + fname)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+
+@pytest.mark.parametrize('prefix', ['', './'])
+def test_extract_dmp_with_leading_dotslash(tmp_path, prefix):
+    # Members named like `./names.dmp` (as produced by GNU tar) must be
+    # found just like plain `names.dmp`. See
+    # https://github.com/etetoolkit/ete/issues/810
+    taxdump = str(tmp_path / 'taxdump.tar.gz')
+    _make_taxdump_tar(taxdump, prefix=prefix)
+
+    dbfile = str(tmp_path / 'taxa.sqlite')
+    ncbiquery.update_db(dbfile, taxdump)  # should not raise
+
+    ncbi = NCBITaxa(dbfile=dbfile)
+    assert ncbi.get_taxid_translator(['2'])[2] == 'Bacteria'
+
+
 def test_ignore_unclassified():
     # normal case
     tree = PhyloTree('((9606, 9598), 10090);')