Merge pull request #70 from rbaldwin-bugseq/add_binding_for_ncbi_output

added binding for saving Taxonomy to ncbi format

Author: audy

src/formats/ncbi.rs

@@ -95,18 +95,24 @@ where
     let mut node_writer = BufWriter::new(std::fs::File::create(dir.join(NODES_FILENAME))?);
     let mut name_writer = BufWriter::new(std::fs::File::create(dir.join(NAMES_FILENAME))?);
 
-    for key in tax.traverse(tax.root())?.filter(|x| x.1).map(|x| x.0) {
+    let root = tax.root();
+    for key in tax.traverse(root.clone())?.filter(|x| x.1).map(|x| x.0) {
         let name = tax.name(key.clone())?;
         let rank = tax.rank(key.clone())?;
+        let parent = if key == root {
+            format!("{}", key)
+        } else {
+            tax.parent(key.clone())?
+                .map(|(x, _)| format!("{}", x))
+                .unwrap_or_default()
+        };
         name_writer
-            .write_all(format!("{}\t|\t{}\t|\tscientific name\t|", &key, name).as_bytes())?;
+            .write_all(format!("{}\t|\t{}\t|\t\t|\tscientific name\t|\n", &key, name).as_bytes())?;
         node_writer.write_all(
             format!(
-                "{}\t|\t{}\t|\t{}\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|",
+                "{}\t|\t{}\t|\t{}\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\t\t|\n",
                 &key,
-                tax.parent(key.clone())?
-                    .map(|(x, _)| format!("{}", x))
-                    .unwrap_or_default(),
+                parent,
                 rank.to_ncbi_rank(),
             )
             .as_bytes(),
@@ -215,5 +221,42 @@ mod tests {
 
         let out = path.join("out");
         save::<&str, _, _>(&tax, &out).unwrap();
+
+        // now load again and validate a few taxids
+        let tax2 = load(&out).unwrap();
+
+        // Check E. coli (562)
+        assert_eq!(
+            Taxonomy::<&str>::name(&tax2, "562").unwrap(),
+            "Escherichia coli"
+        );
+        assert_eq!(
+            Taxonomy::<&str>::rank(&tax2, "562").unwrap(),
+            TaxRank::Species
+        );
+        assert_eq!(
+            Taxonomy::<&str>::parent(&tax2, "562").unwrap(),
+            Some(("561", 1.))
+        );
+
+        // Check Escherichia (561)
+        assert_eq!(Taxonomy::<&str>::name(&tax2, "561").unwrap(), "Escherichia");
+        assert_eq!(
+            Taxonomy::<&str>::rank(&tax2, "561").unwrap(),
+            TaxRank::Genus
+        );
+        assert_eq!(
+            Taxonomy::<&str>::parent(&tax2, "561").unwrap(),
+            Some(("543", 1.))
+        );
+
+        // Check root (1)
+        assert_eq!(Taxonomy::<&str>::name(&tax2, "1").unwrap(), "root");
+
+        // Check children relationship preserved
+        assert_eq!(
+            Taxonomy::<&str>::children(&tax2, "561").unwrap(),
+            vec!["562"]
+        );
     }
 }

src/python.rs

@@ -263,6 +263,19 @@ impl Taxonomy {
         Ok(PyBytes::new(py, &bytes).into())
     }
 
+    /// to_ncbi(self, output_dir: str)
+    /// --
+    ///
+    /// Export a Taxonomy to NCBI format files (nodes.dmp and names.dmp).
+    /// The output directory will be created if it doesn't exist.
+    ///
+    /// Args:
+    ///     output_dir: Path to the directory where nodes.dmp and names.dmp will be written
+    fn to_ncbi(&self, output_dir: &str) -> PyResult<()> {
+        py_try!(ncbi::save::<&str, _, _>(&self.tax, output_dir));
+        Ok(())
+    }
+
     /// to_newick(self)
     /// --
     ///

taxonomy.pyi

@@ -75,6 +75,13 @@ class Taxonomy:
         """Export a Taxonomy as a Newick-encoded byte string."""
         ...
 
+    def to_ncbi(self, output_dir: str) -> None:
+        """
+        Export a Taxonomy to NCBI format files (nodes.dmp and names.dmp).
+        The output directory will be created if it doesn't exist.
+        """
+        ...
+
     def node(self, tax_id: str) -> Optional[TaxonomyNode]:
         """Find a node by its id. Returns `None` if not found"""
         ...