brrrr_lib/
csv_writer.rs

1// (c) Copyright 2020 Trent Hauck
2// All Rights Reserved
3/// The `csv_writer` module provides an implementation for the `RecordWriter` interface to read
4/// and write from csvs.
5use std::io::{self, BufRead, ErrorKind, Write};
6
7use serde::ser::Serialize;
8
9use crate::errors::BrrrrError;
10use crate::types::{FastaRecord, FastqRecord};
11use crate::writer;
12
13use writer::RecordWriter;
14
15use noodles::fasta;
16use noodles::fastq;
17
18/// CsvRecordWriter holds a writer, and outputs FASTA records as newline delimited json.
19pub struct CsvRecordWriter<W: Write> {
20    csv_writer: csv::Writer<W>,
21}
22
23impl<W: Write> CsvRecordWriter<W> {
24    /// Creates a new CsvRecordWriter with a writer.
25    pub fn new(w: W) -> Self {
26        let csv_writer = csv::Writer::from_writer(w);
27        Self { csv_writer }
28    }
29}
30
31impl<W: Write> writer::RecordWriter for CsvRecordWriter<W> {
32    /// Writes an input serializable object to the underlying writer.
33    fn write_serde_record<S: Serialize>(&mut self, r: S) -> io::Result<()> {
34        self.csv_writer.serialize(r)?;
35        Ok(())
36    }
37}
38
39/// Converts a FASTA to CSV
40///
41/// # Arguments
42///
43/// * `input` an input that implements the Read trait.
44/// * `output` an output that implements the Write trait.
45pub fn fa2csv<R: BufRead, W: Write>(input: R, output: &mut W) -> Result<(), BrrrrError> {
46    let mut reader = fasta::Reader::new(input);
47    let record_writer = &mut CsvRecordWriter::new(output);
48
49    for read_record in reader.records() {
50        let record = read_record?;
51        let write_op = record_writer.write_serde_record(FastaRecord::from(record));
52
53        if let Err(e) = write_op {
54            match e.kind() {
55                ErrorKind::BrokenPipe => break,
56                _ => return Err(BrrrrError::IOError(e)),
57            }
58        }
59    }
60    Ok(())
61}
62
63/// Converts a FASTQ file to CSV
64///
65/// # Arguments
66///
67/// * `input` an input that implements the BufRead trait.
68/// * `output` an output that implements the Write trait.
69pub fn fq2csv<R: BufRead, W: Write>(input: R, output: &mut W) -> Result<(), BrrrrError> {
70    let mut reader = fastq::Reader::new(input);
71    let record_writer = &mut CsvRecordWriter::new(output);
72
73    for read_record in reader.records() {
74        let record = read_record?;
75        let write_op = record_writer.write_serde_record(FastqRecord::from(record));
76
77        if let Err(e) = write_op {
78            match e.kind() {
79                ErrorKind::BrokenPipe => break,
80                _ => return Err(BrrrrError::IOError(e)),
81            }
82        }
83    }
84    Ok(())
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_fa2csv() {
93        let input = b">A\nATCG\n" as &[u8];
94
95        let mut output = Vec::new();
96        fa2csv(input, &mut output).unwrap();
97
98        let output_str = String::from_utf8(output).unwrap();
99        let expected_output = "id,desc,seq\nA,,ATCG\n".to_string();
100        assert_eq!(output_str, expected_output);
101    }
102}