Add a method to prune trees back to a rank

onecodex/lib/taxonomy.py

@@ -61,3 +61,19 @@ def generate_skbio_tree(classification, existing_tree=None):
     assert len(unlinked) == 0, 'some unlinked nodes were not included in the tree'
 
     return tree
+
+
+def prune_to_rank(tree, rank='species'):
+    """
+    Takes a TreeNode tree and prunes off any tips not at the specified rank
+    (and backwards up until all of the tips are at the specified rank.
+    """
+    for node in tree.postorder():
+        if node.rank == rank:
+            node._above_rank = True
+        elif any([getattr(n, '_above_rank', False) for n in node.children]):
+            node._above_rank = True
+        else:
+            node._above_rank = False
+
+    tree.remove_deleted(lambda n: not getattr(n, '_above_rank', False))

tests/test_taxonomy.py

@@ -1,4 +1,6 @@
-from onecodex.lib.taxonomy import generate_skbio_tree
+import pytest
+
+from onecodex.lib.taxonomy import generate_skbio_tree, prune_to_rank
 
 
 def test_tree_generation(ocx, api_data):
@@ -8,3 +10,21 @@ def test_tree_generation(ocx, api_data):
     assert tree.is_root()
     staph = tree.find('1279')
     assert staph.tax_name == 'Staphylococcus'
+
+
+def test_tree_pruning(ocx, api_data):
+    from skbio.tree import MissingNodeError
+    classification = ocx.Classifications.get('f9e4a5506b154953')
+    tree = generate_skbio_tree(classification)
+
+    # the species and genus nodes exist
+    tree.find('1279')
+    tree.find('1078083')
+
+    # prune back to genus
+    prune_to_rank(tree, 'genus')
+
+    # the species node no longer exists, but the genus still does
+    with pytest.raises(MissingNodeError):
+        tree.find('1078083')
+    tree.find('1279')