Merge pull request #730 from etetoolkit/ete4_ncbi_prop_get

fix ncbi annotation using other properties
etetoolkit · Nov 8, 2023 · b49543a · b49543a
2 parents 92f28c7 + 10712af
commit b49543a
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/ete4/ncbi_taxonomy/ncbiquery.py b/ete4/ncbi_taxonomy/ncbiquery.py
@@ -475,7 +475,7 @@ def annotate_tree(self, t, taxid_attr="name", tax2name=None,
         taxids = set()
         for n in t.traverse():
             try:
-                tid = int(getattr(n, taxid_attr))
+                tid = int(getattr(n, taxid_attr, n.props.get(taxid_attr)))
             except (ValueError, AttributeError, TypeError):
                 pass
             else:
@@ -501,21 +501,21 @@ def annotate_tree(self, t, taxid_attr="name", tax2name=None,
 
         for n in t.traverse('postorder'):
             try:
-                node_taxid = int(getattr(n, taxid_attr))
+                node_taxid = int(getattr(n, taxid_attr, n.props.get(taxid_attr)))
             except (ValueError, AttributeError, TypeError):
                 node_taxid = None
 
             n.add_prop('taxid', node_taxid)
             if node_taxid:
                 if node_taxid in merged_conversion:
                     node_taxid = merged_conversion[node_taxid]
-                n.add_props(sci_name = tax2name.get(node_taxid, getattr(n, taxid_attr, '')),
+                n.add_props(sci_name = tax2name.get(node_taxid, getattr(n, taxid_attr, n.props.get(taxid_attr, ''))),
                                common_name = tax2common_name.get(node_taxid, ''),
                                lineage = tax2track.get(node_taxid, []),
                                rank = tax2rank.get(node_taxid, 'Unknown'),
                                named_lineage = [tax2name.get(tax, str(tax)) for tax in tax2track.get(node_taxid, [])])
             elif n.is_leaf:
-                n.add_props(sci_name = getattr(n, taxid_attr, 'NA'),
+                n.add_props(sci_name = getattr(n, taxid_attr, n.props.get(taxid_attr, 'NA')),
                                common_name = '',
                                lineage = [],
                                rank = 'Unknown',