diff --git a/ete4/ncbi_taxonomy/ncbiquery.py b/ete4/ncbi_taxonomy/ncbiquery.py index 5048003bb..eebc2e86b 100644 --- a/ete4/ncbi_taxonomy/ncbiquery.py +++ b/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') diff --git a/tests/test_ncbiquery.py b/tests/test_ncbiquery.py index 3663a756d..e0ebf6f42 100644 --- a/tests/test_ncbiquery.py +++ b/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);')