Add a method to get all descendants of a node

Author: Keats

Cargo.toml

@@ -12,7 +12,7 @@ edition = "2021"
 
 [dependencies]
 memchr = "2.2.1"
-pyo3 = { version = "0.17", optional = true }
+pyo3 = { version = "0.18", optional = true }
 quick-xml = "0.27"
 serde = { version = "1.0.104", features = ["derive"] }
 serde_json = "1.0.44"

README.md

@@ -110,7 +110,12 @@ In NCBI, it only accounts for *scientific names* and not synonyms.
 
 #### `tax.children(tax_id: str) -> List[TaxonomyNode]`
 
-Returns all nodes below the given tax id.
+Returns all direct nodes below the given tax id.
+
+#### `tax.descendants(tax_id: str) -> List[TaxonomyNode]`
+
+Returns all nodes below the given tax id. 
+Equivalent to running `tax.children` recursively on the initial result of `tax.children(tax_id)`.
 
 #### `tax.lineage(tax_id: str) -> List[TaxonomyNode]`
 

src/base.rs

@@ -324,6 +324,17 @@ impl<'t> Taxonomy<'t, &'t str> for GeneralTaxonomy {
             .collect()
     }
 
+    fn descendants(&'t self, tax_id: &'t str) -> TaxonomyResult<Vec<&'t str>> {
+        let children: HashSet<&str> = self
+            .traverse(tax_id)?
+            .map(|(n, _)| n)
+            .filter(|t| *t != tax_id)
+            .collect();
+        let mut children: Vec<&str> = children.into_iter().collect();
+        children.sort_unstable();
+        Ok(children)
+    }
+
     fn parent(&'t self, tax_id: &str) -> TaxonomyResult<Option<(&'t str, f32)>> {
         let idx = self.to_internal_index(tax_id)?;
         if idx == 0 {
@@ -377,6 +388,15 @@ impl<'t> Taxonomy<'t, InternalIndex> for GeneralTaxonomy {
         }
     }
 
+    fn descendants(&'t self, tax_id: InternalIndex) -> TaxonomyResult<Vec<InternalIndex>> {
+        let children: HashSet<InternalIndex> = self
+            .traverse(tax_id)?
+            .map(|(n, _)| n)
+            .filter(|n| *n != tax_id)
+            .collect();
+        Ok(children.into_iter().collect())
+    }
+
     fn parent(&'t self, idx: InternalIndex) -> TaxonomyResult<Option<(InternalIndex, f32)>> {
         if idx == 0 {
             return Ok(None);
@@ -452,6 +472,10 @@ mod tests {
         let tax = create_test_taxonomy();
         assert_eq!(Taxonomy::<&str>::len(&tax), 6);
         assert_eq!(tax.children("1").unwrap(), vec!["2", "1000", "101", "102"]);
+        assert_eq!(
+            tax.descendants("1").unwrap(),
+            vec!["1000", "101", "102", "2", "562"]
+        );
         assert_eq!(tax.name("562").unwrap(), "Escherichia coli");
         assert_eq!(tax.rank("562").unwrap(), TaxRank::Species);
         assert_eq!(tax.parent("562").unwrap(), Some(("2", 1.0)));

src/python.rs

@@ -290,12 +290,13 @@ impl Taxonomy {
     ///
     /// Find a node by its name, Raises an exception if not found.
     fn find_all_by_name(&self, name: &str) -> PyResult<Vec<TaxonomyNode>> {
-        Ok(self
+        let res = self
             .tax
             .find_all_by_name(name)
-            .iter()
-            .map(|tax_id| self.as_node(tax_id).unwrap())
-            .collect::<Vec<TaxonomyNode>>())
+            .into_iter()
+            .map(|tax_id| self.as_node(tax_id))
+            .collect::<PyResult<Vec<TaxonomyNode>>>()?;
+        Ok(res)
     }
 
     /// parent_with_distance(self, tax_id: str, /, at_rank: str)
@@ -345,22 +346,24 @@ impl Taxonomy {
     /// children(self, tax_id: str)
     /// --
     ///
-    /// Return a list of child taxonomy nodes from the node id provided.
+    /// Return a list of direct child taxonomy nodes from the node id provided.
     fn children(&self, tax_id: &str) -> PyResult<Vec<TaxonomyNode>> {
-        let children: Vec<&str> = py_try!(self.tax.children(tax_id));
-        let mut res = Vec::with_capacity(children.len());
-        for key in children {
-            let child = self.node(key);
-            if let Some(c) = child {
-                res.push(c);
-            } else {
-                return Err(PyErr::new::<TaxonomyError, _>(format!(
-                    "Node {} is missing in children",
-                    key
-                )));
-            }
-        }
+        let res = py_try!(self.tax.children(tax_id))
+            .into_iter()
+            .map(|tax_id| self.as_node(tax_id))
+            .collect::<PyResult<Vec<TaxonomyNode>>>()?;
+        Ok(res)
+    }
 
+    /// descendants(self, tax_id: str)
+    /// --
+    ///
+    /// Return a list of all child taxonomy nodes from the node id provided.
+    fn descendants(&self, tax_id: &str) -> PyResult<Vec<TaxonomyNode>> {
+        let res = py_try!(self.tax.descendants(tax_id))
+            .into_iter()
+            .map(|tax_id| self.as_node(tax_id))
+            .collect::<PyResult<Vec<TaxonomyNode>>>()?;
         Ok(res)
     }
 
@@ -370,20 +373,10 @@ impl Taxonomy {
     /// Return a list of all the parent taxonomy nodes of the node id provided
     /// (including that node itself).
     fn lineage(&self, tax_id: &str) -> PyResult<Vec<TaxonomyNode>> {
-        let lineage: Vec<&str> = py_try!(self.tax.lineage(tax_id));
-        let mut res = Vec::with_capacity(lineage.len());
-        for key in lineage {
-            let ancestor = self.node(key);
-            if let Some(a) = ancestor {
-                res.push(a);
-            } else {
-                return Err(PyErr::new::<TaxonomyError, _>(format!(
-                    "Node {} is missing in lineage",
-                    key
-                )));
-            }
-        }
-
+        let res = py_try!(self.tax.lineage(tax_id))
+            .into_iter()
+            .map(|tax_id| self.as_node(tax_id))
+            .collect::<PyResult<Vec<TaxonomyNode>>>()?;
         Ok(res)
     }
 

src/taxonomy.rs

@@ -15,9 +15,12 @@ where
     /// Returns the root node of the entire tree.
     fn root(&'t self) -> T;
 
-    /// Returns a [Vec] of all the child IDs of the given tax_id.
+    /// Returns a [Vec] of all the direct children IDs of the given tax_id.
     fn children(&'t self, tax_id: T) -> TaxonomyResult<Vec<T>>;
 
+    /// Returns a [Vec] of all the children IDs of the given tax_id.
+    fn descendants(&'t self, tax_id: T) -> TaxonomyResult<Vec<T>>;
+
     /// Returns the parent of the given taxonomic node and the distance to said parent.
     /// The parent of the root node will return [None].
     fn parent(&'t self, tax_id: T) -> TaxonomyResult<Option<(T, f32)>>;
@@ -208,6 +211,17 @@ pub(crate) mod tests {
             })
         }
 
+        fn descendants(&'t self, tax_id: u32) -> TaxonomyResult<Vec<u32>> {
+            let children: HashSet<u32> = self
+                .traverse(tax_id)?
+                .map(|(n, _)| n)
+                .filter(|n| *n != tax_id)
+                .collect();
+            let mut children: Vec<u32> = children.into_iter().collect();
+            children.sort_unstable();
+            Ok(children)
+        }
+
         fn parent(&self, tax_id: u32) -> TaxonomyResult<Option<(u32, f32)>> {
             Ok(match tax_id {
                 131567 => Some((1, 1.)),
@@ -265,6 +279,15 @@ pub(crate) mod tests {
         assert_eq!(tax.is_empty(), false);
     }
 
+    #[test]
+    fn test_descendants() {
+        let tax = MockTax;
+        assert_eq!(
+            tax.descendants(2).unwrap(),
+            vec![22, 1046, 1224, 1236, 53452, 56812, 61598, 62322, 135613, 135622, 765909]
+        );
+    }
+
     #[test]
     fn test_lca() {
         let tax = MockTax;