fastq_stats/

lib.rs

//! # Fastq file statistics
//! This is a minimal Python package to quickly count reads and basepairs in `.fastq.gz` files.
//!
//! To achieve its performance goals, the package is implemented in Rust and packages for Python
//! using the [Maturin project](https://www.maturin.rs/).
//!
//! The package contains two functions: [`analyze_file()`] and [`analyze_files()`].
//! Both return an instance of [`ReadsStats`], which exposes the [`read_count`](ReadsStats::read_count)
//! and [`basepair_count`](ReadsStats::basepair_count) of the input files.
//!
//! # Python Usage Example
//!
//! ```python
//! from pathlib import Path
//! from fastq_stats import ReadsStats, analyze_file, analyze_files
//!
//! # single thread
//! result = analyze_file(Path("file.fastq.gz"))
//! assert result == ReadsStats(read_count=10, basepair_count=120)
//!
//! # multiple threads (default: up to os.cpu_count())
//! result = analyze_files([Path("file1.fastq.gz"), Path("file2.fastq.gz")])
//! assert result == ReadsStats(read_count=12, basepair_count=250)
//!
//! # multiple threads (explicit)
//! result = analyze_files([Path("file1.fastq.gz"), Path("file2.fastq.gz")], num_threads=3)
//! assert result == ReadsStats(read_count=12, basepair_count=250)
//! ```

use flate2::read::MultiGzDecoder;
use pyo3::prelude::*;
use rayon::prelude::*;
use rayon::ThreadPoolBuilder;
use std::fmt;
use std::fmt::{Display, Formatter};
use std::fs::File;
use std::io::{BufRead, BufReader, Error};
use std::iter::Sum;
use std::ops::Add;
use std::path::PathBuf;

/// Statistics about NGS reads
#[pyclass(eq, frozen)]
#[derive(Debug, PartialEq)]
pub struct ReadsStats {
    /// Number of reads
    #[pyo3(get)]
    pub read_count: u64,
    /// Number of basepairs
    #[pyo3(get)]
    pub basepair_count: u64,
}

impl Display for ReadsStats {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        // Write strictly the first element into the supplied output
        // stream: `f`. Returns `fmt::Result` which indicates whether the
        // operation succeeded or failed. Note that `write!` uses syntax which
        // is very similar to `println!`.
        write!(
            f,
            "ReadsStats(read_count={}, basepair_count={})",
            self.read_count, self.basepair_count
        )
    }
}

impl Add for ReadsStats {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            read_count: self.read_count + other.read_count,
            basepair_count: self.basepair_count + other.basepair_count,
        }
    }
}

impl Sum for ReadsStats {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(
            Self {
                read_count: 0,
                basepair_count: 0,
            },
            |a, b| a + b,
        )
    }
}

#[pymethods]
impl ReadsStats {
    #[new]
    fn new(read_count: u64, basepair_count: u64) -> Self {
        Self {
            read_count,
            basepair_count,
        }
    }

    fn __repr__(&self) -> String {
        format!("{self}")
    }

    fn __add__(&self, other: &Self) -> Self {
        ReadsStats {
            read_count: self.read_count + other.read_count,
            basepair_count: self.basepair_count + other.basepair_count,
        }
    }
}

/// Collect `.fastq.gz` file statistics
#[pymodule]
fn fastq_stats(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<ReadsStats>()?;
    module.add_function(wrap_pyfunction!(analyze_file, module)?)?;
    module.add_function(wrap_pyfunction!(analyze_files, module)?)?;
    Ok(())
}

/// Collection of Path-like Python types
#[derive(FromPyObject, Clone)]
pub enum PathLike {
    String(String),
    Path(PathBuf),
}

/// Count all reads and basepairs in multiple `.fastq.gz` files.
///
/// # Errors
///
/// Will return `Err` if a file could not be read or decoded.
///
/// # Panics
///
/// Will panic if the requested `ThreadPool` could not be built.
#[pyfunction]
#[pyo3(signature = (filepaths, *, num_threads=None))]
#[allow(clippy::needless_pass_by_value)] // pyfunction requires value types
pub fn analyze_files(
    filepaths: Vec<PathLike>,
    num_threads: Option<usize>,
) -> Result<ReadsStats, Error> {
    let task = || {
        filepaths
            .par_iter()
            .map(|filepath| analyze_file(filepath.clone()))
            .sum()
    };

    match num_threads {
        None => task(), // execute in default threadpool
        Some(threads) => {
            let pool = ThreadPoolBuilder::new()
                .num_threads(threads)
                .build()
                .expect("Failed to build thread pool");
            pool.install(task)
        }
    }
}

/// Count all reads and basepairs in a `.fastq.gz` file.
///
/// # Errors
///
/// Will return `Err` if the file could not be read or decoded.
#[pyfunction]
pub fn analyze_file(filepath: PathLike) -> Result<ReadsStats, Error> {
    let file = match filepath {
        PathLike::String(path) => File::open(path),
        PathLike::Path(path) => File::open(path),
    }?;
    let gz_reader = MultiGzDecoder::new(file);
    let buf_gz_reader = BufReader::new(gz_reader);

    let mut reads: u64 = 0;
    let mut basepairs: u64 = 0;
    for (linenum, line) in buf_gz_reader.lines().enumerate() {
        let line = line?;
        if linenum % 4 != 1 {
            continue;
        }
        reads += 1;
        basepairs += line.len() as u64;
    }
    Ok(ReadsStats {
        read_count: reads,
        basepair_count: basepairs,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use std::fs;
    use std::io::Write;

    struct TempFile {
        pub wrapped_path: PathBuf,
    }

    impl Drop for TempFile {
        fn drop(&mut self) {
            if let Ok(()) = fs::remove_file(&self.wrapped_path) {};
        }
    }

    impl From<PathBuf> for TempFile {
        fn from(path: PathBuf) -> Self {
            Self { wrapped_path: path }
        }
    }

    impl From<&str> for TempFile {
        fn from(path: &str) -> Self {
            Self {
                wrapped_path: PathBuf::from(path),
            }
        }
    }

    #[test]
    fn test_single_read_single_gz() {
        // prepare test file
        let f = TempFile::from("testfile1.fastq.gz");
        let mut buf = Vec::new();

        // write gzipped sequence
        let mut encoder = GzEncoder::new(&mut buf, Compression::default());
        encoder.write_all(b"@seq1\nACTG\n+\nooaa\n").unwrap();
        encoder.flush().unwrap();
        encoder.finish().unwrap();
        fs::write(&f.wrapped_path, buf).unwrap();

        // test
        let result = analyze_file(PathLike::Path(f.wrapped_path.clone()).clone()).unwrap();
        assert_eq!(result.read_count, 1);
        assert_eq!(result.basepair_count, 4);
    }

    #[test]
    fn test_multi_read_single_gz() {
        // prepare test file
        let f = TempFile::from("testfile2.fastq.gz");
        let mut buf = Vec::new();

        // write gzipped sequence
        let mut encoder = GzEncoder::new(&mut buf, Compression::default());
        encoder
            .write_all(b"@seq1\nACTG\n+\nooaa\n@seq2\nAAA\n+\naao\n")
            .unwrap();
        encoder.flush().unwrap();
        encoder.finish().unwrap();
        fs::write(&f.wrapped_path, buf).unwrap();

        // test
        let result = analyze_file(PathLike::Path(f.wrapped_path.clone()).clone()).unwrap();
        assert_eq!(result.read_count, 2);
        assert_eq!(result.basepair_count, 7);
    }

    #[test]
    fn test_multi_read_multi_gz() {
        // prepare test file
        let f = TempFile::from("testfile3.fastq.gz");
        let mut buf1 = Vec::new();

        // write first gzipped sequence
        let mut encoder = GzEncoder::new(&mut buf1, Compression::default());
        encoder.write_all(b"@seq1\nACTG\n+\nooaa\n").unwrap();
        encoder.flush().unwrap();
        encoder.finish().unwrap();

        // write second gzipped sequence
        let mut encoder = GzEncoder::new(&mut buf1, Compression::default());
        encoder.write_all(b"@seq2\nACA\n+\nooa\n").unwrap();
        encoder.flush().unwrap();
        encoder.finish().unwrap();
        fs::write(&f.wrapped_path, buf1).unwrap();

        // test
        let result = analyze_file(PathLike::Path(f.wrapped_path.clone()).clone()).unwrap();
        assert_eq!(result.read_count, 2);
        assert_eq!(result.basepair_count, 7);
    }
}