Some refactoring in readwrite module

Author: sirno

src/errors.rs

@@ -2,10 +2,13 @@
 
 use std::fmt;
 
+pub type Result<T> = std::result::Result<T, VirolutionError>;
+
 #[derive(Clone, Debug)]
 pub enum VirolutionError {
     ImplementationError(String),
     InitializationError(String),
+    ReadError(String),
 }
 
 impl fmt::Display for VirolutionError {
@@ -17,6 +20,9 @@ impl fmt::Display for VirolutionError {
             VirolutionError::InitializationError(message) => {
                 write!(f, "InitializationError: {}", message)
             }
+            VirolutionError::ReadError(message) => {
+                write!(f, "ReadError: {}", message)
+            }
         }
     }
 }

src/readwrite/haplotype.rs

@@ -1,28 +1,38 @@
+use seq_io::fasta;
+use seq_io::fasta::Record;
+
 use crate::core::haplotype::{Haplotype, Symbol, Wildtype};
 use crate::encoding::STRICT_DECODE;
+use crate::errors::{Result, VirolutionError};
 use crate::references::HaplotypeRef;
-use seq_io::fasta;
-use seq_io::fasta::Record;
 
 pub trait HaplotypeIO {
-    fn load_wildtype(path: &str) -> Result<HaplotypeRef, fasta::Error>;
+    fn load_wildtype(path: &str) -> Result<HaplotypeRef>;
 }
 
 impl HaplotypeIO for Haplotype {
-    fn load_wildtype(path: &str) -> Result<HaplotypeRef, fasta::Error> {
-        let mut reader = fasta::Reader::from_path(path)?;
-        let sequence: Vec<Symbol> = reader
-            .next()
-            .unwrap()
-            .expect("Unable to read sequence.")
-            .seq()
-            .iter()
-            .filter(|&&enc| enc != 0x0au8 && enc != 0x0du8)
-            .map(|enc| match STRICT_DECODE.get(enc) {
-                Some(result) => Some(*result),
-                None => panic!("Unable to decode literal {enc}."),
-            })
-            .collect();
-        Ok(Wildtype::new(sequence))
+    fn load_wildtype(path: &str) -> Result<HaplotypeRef> {
+        let mut reader = fasta::Reader::from_path(path).map_err(|_| {
+            VirolutionError::InitializationError(format!(
+                "Unable create file reader for fasta file: {path}"
+            ))
+        })?;
+
+        match reader.next() {
+            None => Err(VirolutionError::InitializationError(format!(
+                "No sequence found in fasta file: {path}"
+            ))),
+            Some(Err(_)) => Err(VirolutionError::InitializationError(format!(
+                "Unable to read sequence from fasta file: {path}"
+            ))),
+            Some(Ok(sequence_record)) => {
+                let sequence: Vec<Symbol> = sequence_record
+                    .seq()
+                    .iter()
+                    .filter_map(|enc| STRICT_DECODE.get(enc).copied().map(Some))
+                    .collect();
+                Ok(Wildtype::new(sequence))
+            }
+        }
     }
 }

src/readwrite/population.rs

@@ -1,11 +1,13 @@
+use serde::Deserialize;
+
 use crate::core::haplotype::Symbol;
 use crate::core::population::Population;
 use crate::encoding::DECODE;
+use crate::errors::{Result, VirolutionError};
 use crate::references::HaplotypeRef;
-use serde::Deserialize;
 
 pub trait PopulationIO {
-    fn read(path: &str, wildtype: HaplotypeRef) -> Result<Population, csv::Error>;
+    fn read(path: &str, wildtype: HaplotypeRef) -> Result<Population>;
     fn write(&self, path: &str);
 }
 
@@ -20,30 +22,26 @@ impl PopulationIO for Population {
     ///
     /// Warning: This function will create multiple instances of the same
     /// haplotype if it is present multiple times in the CSV file.
-    fn read(path: &str, wildtype: HaplotypeRef) -> Result<Population, csv::Error> {
-        let mut reader = csv::Reader::from_path(path)?;
+    fn read(path: &str, wildtype: HaplotypeRef) -> Result<Population> {
+        let mut reader = csv::Reader::from_path(path)
+            .map_err(|_err| VirolutionError::ReadError(format!("Failed to read from {path}")))?;
         let mut populations: Vec<Population> = Vec::new();
 
         for record in reader.deserialize() {
-            let record: HaplotypeRecord = record?;
+            let record: HaplotypeRecord = record.map_err(|_err| {
+                VirolutionError::ReadError(format!("Failed to parse record in {path}"))
+            })?;
+
+            // skip wildtype
             if record.haplotype == "wt" {
                 continue;
             }
-            let mutations = record.haplotype.split(';');
-            let mut positions: Vec<usize> = Vec::new();
-            let mut changes: Vec<Symbol> = Vec::new();
-            mutations.for_each(|mutation| {
-                let mut mutation = mutation.split(':');
-                let position = mutation.next().unwrap();
-
-                let mut change = mutation.next().unwrap().split("->");
-                let _origin = change.next();
-                let target = change.next().unwrap().chars().next().unwrap() as u8;
-
-                positions.push(position.parse::<usize>().unwrap());
-                changes.push(DECODE.get(&target).copied());
-            });
+
+            // parse haplotype
+            let (positions, changes) = parse_haplotype(&record.haplotype)?;
             let haplotype = wildtype.create_descendant(positions, changes, 0);
+
+            // create population and add for merging
             let population = Population::from_haplotype(haplotype, record.count);
             populations.push(population);
         }
@@ -55,3 +53,37 @@ impl PopulationIO for Population {
         unimplemented!()
     }
 }
+
+fn parse_haplotype(haplotype: &str) -> Result<(Vec<usize>, Vec<Symbol>)> {
+    let mutations = haplotype.split(';');
+    let parsed_mutations: Result<Vec<(usize, Symbol)>> = mutations.map(parse_mutation).collect();
+    Ok(parsed_mutations?.into_iter().unzip())
+}
+
+fn parse_mutation(mutation: &str) -> Result<(usize, Symbol)> {
+    let mutation_split: Vec<&str> = mutation.split(':').collect();
+    match mutation_split.as_slice() {
+        [position, change] => Ok((
+            position.parse::<usize>().map_err(|_| {
+                VirolutionError::ReadError(format!("Failed to parse position: {position}"))
+            })?,
+            parse_change(change)?,
+        )),
+        _ => Err(VirolutionError::ReadError(format!(
+            "Invalid mutation format: {mutation}"
+        ))),
+    }
+}
+
+fn parse_change(change: &str) -> Result<Symbol> {
+    let change_split: Vec<&str> = change.split(':').collect();
+    match change_split.as_slice() {
+        [_origin, target] => {
+            let target = target.chars().next().unwrap() as u8;
+            Ok(DECODE.get(&target).copied())
+        }
+        _ => Err(VirolutionError::ReadError(format!(
+            "Invalid change format: {change}"
+        ))),
+    }
+}