fastx/
writer.rs

1//! FASTA/FASTQ writer with optional gzip compression and validation.
2
3use std::fmt;
4use std::fs::File;
5use std::io::{self, BufWriter, Write};
6use std::path::Path;
7
8use crate::error::{Error, Result};
9use crate::format::{Compression, CompressionLevel, Format};
10use crate::record::Sequence;
11use crate::seq::Alphabet;
12
13/// Default FASTA line width. NCBI uses 70, UCSC and samtools use 60.
14pub const DEFAULT_LINE_WIDTH: usize = 60;
15
16/// A writer for FASTA and FASTQ records.
17///
18/// FASTA sequence lines are wrapped at [`DEFAULT_LINE_WIDTH`] unless configured
19/// otherwise. Writing a FASTQ record without quality scores is an error rather
20/// than a silent fabrication of quality values; writing a FASTQ record to a FASTA
21/// writer simply drops the quality, which is the usual `fq2fa` conversion.
22///
23/// ```
24/// use fastx::{FastxWriter, Format, Sequence};
25///
26/// let mut out = Vec::new();
27/// {
28///     let mut writer = FastxWriter::new(&mut out, Format::Fasta).line_width(4);
29///     writer.write_record(&Sequence::fasta("a", b"ACGTAC"))?;
30///     writer.flush()?;
31/// }
32/// assert_eq!(out, b">a\nACGT\nAC\n");
33/// # Ok::<(), fastx::Error>(())
34/// ```
35pub struct FastxWriter<W: Write> {
36    inner: W,
37    format: Format,
38    line_width: Option<usize>,
39    validate: Option<Alphabet>,
40    written: u64,
41}
42
43impl<W: Write> FastxWriter<W> {
44    /// A writer for `format`, wrapping FASTA at [`DEFAULT_LINE_WIDTH`].
45    ///
46    /// Wrap `inner` in a [`BufWriter`] yourself for best throughput, or use
47    /// [`create`] which does it for you.
48    pub fn new(inner: W, format: Format) -> FastxWriter<W> {
49        FastxWriter {
50            inner,
51            format,
52            line_width: Some(DEFAULT_LINE_WIDTH),
53            validate: None,
54            written: 0,
55        }
56    }
57
58    /// Set the FASTA line width. `0` disables wrapping.
59    pub fn line_width(mut self, width: usize) -> Self {
60        self.line_width = if width == 0 { None } else { Some(width) };
61        self
62    }
63
64    /// Validate every record against `alphabet` before writing it.
65    pub fn validate(mut self, alphabet: Alphabet) -> Self {
66        self.validate = Some(alphabet);
67        self
68    }
69
70    /// The output format.
71    pub fn format(&self) -> Format {
72        self.format
73    }
74
75    /// Number of records written so far.
76    pub fn records_written(&self) -> u64 {
77        self.written
78    }
79
80    /// Write one record in the writer's format.
81    pub fn write_record(&mut self, record: &Sequence) -> Result<()> {
82        if let Some(alphabet) = self.validate {
83            record.validate(alphabet)?;
84        }
85        match self.format {
86            Format::Fasta => record.write_fasta(&mut self.inner, self.line_width)?,
87            Format::Fastq => record.write_fastq(&mut self.inner)?,
88        }
89        self.written += 1;
90        Ok(())
91    }
92
93    /// Write every record of an iterator.
94    pub fn write_all<'a, I>(&mut self, records: I) -> Result<()>
95    where
96        I: IntoIterator<Item = &'a Sequence>,
97    {
98        for record in records {
99            self.write_record(record)?;
100        }
101        Ok(())
102    }
103
104    /// Write a FASTA record from parts, without building a [`Sequence`].
105    pub fn write_fasta(&mut self, id: &str, description: Option<&str>, seq: &[u8]) -> Result<()> {
106        crate::record::check_writable_residues(seq, id)?;
107        self.write_header(b'>', id, description)?;
108        match self.line_width {
109            None => {
110                self.inner.write_all(seq)?;
111                self.inner.write_all(b"\n")?;
112            }
113            Some(width) => {
114                if seq.is_empty() {
115                    self.inner.write_all(b"\n")?;
116                }
117                for chunk in seq.chunks(width) {
118                    self.inner.write_all(chunk)?;
119                    self.inner.write_all(b"\n")?;
120                }
121            }
122        }
123        self.written += 1;
124        Ok(())
125    }
126
127    /// Write a FASTQ record from parts, without building a [`Sequence`].
128    pub fn write_fastq(
129        &mut self,
130        id: &str,
131        description: Option<&str>,
132        seq: &[u8],
133        quality: &[u8],
134    ) -> Result<()> {
135        if seq.len() != quality.len() {
136            return Err(Error::LengthMismatch {
137                id: id.to_string(),
138                seq: seq.len(),
139                quality: quality.len(),
140            });
141        }
142        crate::record::check_writable_residues(seq, id)?;
143        crate::record::check_writable_fastq_sequence(seq, id)?;
144        crate::record::check_writable_quality(quality, id)?;
145        self.write_header(b'@', id, description)?;
146        self.inner.write_all(seq)?;
147        self.inner.write_all(b"\n+\n")?;
148        self.inner.write_all(quality)?;
149        self.inner.write_all(b"\n")?;
150        self.written += 1;
151        Ok(())
152    }
153
154    fn write_header(&mut self, prefix: u8, id: &str, description: Option<&str>) -> Result<()> {
155        self.inner.write_all(&[prefix])?;
156        self.inner.write_all(id.as_bytes())?;
157        if let Some(description) = description.filter(|d| !d.is_empty()) {
158            self.inner.write_all(b" ")?;
159            self.inner.write_all(description.as_bytes())?;
160        }
161        self.inner.write_all(b"\n")?;
162        Ok(())
163    }
164
165    /// Flush buffered bytes to the underlying writer.
166    pub fn flush(&mut self) -> Result<()> {
167        self.inner.flush()?;
168        Ok(())
169    }
170
171    /// Flush and return the underlying writer.
172    ///
173    /// Always prefer this over dropping the writer: for gzip output it is the
174    /// only way to learn that finishing the stream failed.
175    pub fn finish(mut self) -> Result<W> {
176        self.inner.flush()?;
177        Ok(self.inner)
178    }
179
180    /// Borrow the underlying writer.
181    pub fn get_ref(&self) -> &W {
182        &self.inner
183    }
184}
185
186impl<W: Write> fmt::Debug for FastxWriter<W> {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        f.debug_struct("FastxWriter")
189            .field("format", &self.format)
190            .field("line_width", &self.line_width)
191            .field("validate", &self.validate)
192            .field("records_written", &self.written)
193            .finish_non_exhaustive()
194    }
195}
196
197/// Configuration for a [`FastxWriter`].
198///
199/// ```no_run
200/// use fastx::{Alphabet, WriterBuilder};
201///
202/// // Format and gzip are inferred from the path: this writes gzipped FASTQ.
203/// let mut writer = WriterBuilder::new()
204///     .line_width(80)
205///     .validate(Alphabet::Dna)
206///     .create("out.fq.gz")?;
207/// writer.finish()?;
208/// # Ok::<(), fastx::Error>(())
209/// ```
210#[derive(Debug, Clone)]
211pub struct WriterBuilder {
212    format: Option<Format>,
213    line_width: Option<usize>,
214    compression: Option<Compression>,
215    level: CompressionLevel,
216    validate: Option<Alphabet>,
217    buffer_size: usize,
218}
219
220impl Default for WriterBuilder {
221    fn default() -> Self {
222        WriterBuilder {
223            format: None,
224            line_width: Some(DEFAULT_LINE_WIDTH),
225            compression: None,
226            level: CompressionLevel::default(),
227            validate: None,
228            buffer_size: 128 * 1024,
229        }
230    }
231}
232
233impl WriterBuilder {
234    /// A builder with default settings.
235    pub fn new() -> WriterBuilder {
236        WriterBuilder::default()
237    }
238
239    /// Force the output format instead of inferring it from the path.
240    pub fn format(mut self, format: Format) -> Self {
241        self.format = Some(format);
242        self
243    }
244
245    /// FASTA line width; `0` disables wrapping.
246    pub fn line_width(mut self, width: usize) -> Self {
247        self.line_width = if width == 0 { None } else { Some(width) };
248        self
249    }
250
251    /// Force compression instead of inferring it from the path.
252    pub fn compression(mut self, compression: Compression) -> Self {
253        self.compression = Some(compression);
254        self
255    }
256
257    /// gzip compression level.
258    pub fn level(mut self, level: CompressionLevel) -> Self {
259        self.level = level;
260        self
261    }
262
263    /// Validate records against an alphabet before writing.
264    pub fn validate(mut self, alphabet: Alphabet) -> Self {
265        self.validate = Some(alphabet);
266        self
267    }
268
269    /// Output buffer size in bytes.
270    pub fn buffer_size(mut self, bytes: usize) -> Self {
271        self.buffer_size = bytes;
272        self
273    }
274
275    /// Build a writer around any [`Write`]. Requires a format to be set.
276    pub fn build<W: Write>(&self, inner: W) -> Result<FastxWriter<W>> {
277        let format = self.format.ok_or(Error::UnknownFormat {
278            hint: "no format given and no path to infer it from".to_string(),
279        })?;
280        Ok(FastxWriter {
281            inner,
282            format,
283            line_width: self.line_width,
284            validate: self.validate,
285            written: 0,
286        })
287    }
288
289    /// Create a file, inferring format and gzip compression from its extension.
290    pub fn create<P: AsRef<Path>>(&self, path: P) -> Result<BoxedWriter> {
291        let path = path.as_ref();
292        let format = match self.format.or_else(|| Format::from_path(path)) {
293            Some(format) => format,
294            None => {
295                return Err(Error::UnknownFormat {
296                    hint: format!("unrecognised extension in {}", path.display()),
297                })
298            }
299        };
300        // Compressed output defaults to BGZF rather than plain gzip. It is valid
301        // gzip either way, so nothing downstream breaks, but BGZF can later be
302        // indexed and randomly accessed, which plain gzip never can. Ask for
303        // `Compression::Gzip` explicitly to get a single deflate stream.
304        let compression = self
305            .compression
306            .unwrap_or(match Compression::from_path(path) {
307                Compression::Gzip => Compression::Bgzf,
308                other => other,
309            });
310        let file = File::create(path)
311            .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
312        let buffered = BufWriter::with_capacity(self.buffer_size, file);
313        let sink: Box<dyn Write + Send> = match compression {
314            Compression::None => Box::new(buffered),
315            Compression::Gzip => gzip_writer(buffered, self.level)?,
316            Compression::Bgzf => bgzf_writer(buffered, self.level)?,
317        };
318        Ok(FastxWriter {
319            inner: sink,
320            format,
321            line_width: self.line_width,
322            validate: self.validate,
323            written: 0,
324        })
325    }
326
327    /// Write to standard output.
328    pub fn stdout(&self) -> Result<BoxedWriter> {
329        let format = self.format.unwrap_or(Format::Fasta);
330        let sink: Box<dyn Write + Send> = match self.compression.unwrap_or(Compression::None) {
331            Compression::None => Box::new(BufWriter::with_capacity(self.buffer_size, io::stdout())),
332            Compression::Gzip => gzip_writer(
333                BufWriter::with_capacity(self.buffer_size, io::stdout()),
334                self.level,
335            )?,
336            Compression::Bgzf => bgzf_writer(
337                BufWriter::with_capacity(self.buffer_size, io::stdout()),
338                self.level,
339            )?,
340        };
341        Ok(FastxWriter {
342            inner: sink,
343            format,
344            line_width: self.line_width,
345            validate: self.validate,
346            written: 0,
347        })
348    }
349}
350
351#[cfg(feature = "gzip")]
352fn gzip_writer<W: Write + Send + 'static>(
353    sink: W,
354    level: CompressionLevel,
355) -> Result<Box<dyn Write + Send>> {
356    Ok(Box::new(flate2::write::GzEncoder::new(
357        sink,
358        flate2::Compression::new(level.0.min(9)),
359    )))
360}
361
362#[cfg(not(feature = "gzip"))]
363fn gzip_writer<W: Write + Send + 'static>(
364    _sink: W,
365    _level: CompressionLevel,
366) -> Result<Box<dyn Write + Send>> {
367    Err(Error::FeatureDisabled("gzip"))
368}
369
370#[cfg(feature = "gzip")]
371fn bgzf_writer<W: Write + Send + 'static>(
372    sink: W,
373    level: CompressionLevel,
374) -> Result<Box<dyn Write + Send>> {
375    Ok(Box::new(crate::bgzf::BgzfWriter::with_level(sink, level)))
376}
377
378#[cfg(not(feature = "gzip"))]
379fn bgzf_writer<W: Write + Send + 'static>(
380    _sink: W,
381    _level: CompressionLevel,
382) -> Result<Box<dyn Write + Send>> {
383    Err(Error::FeatureDisabled("gzip"))
384}
385
386/// The boxed writer type produced by [`create`] and [`WriterBuilder::stdout`].
387pub type BoxedWriter = FastxWriter<Box<dyn Write + Send>>;
388
389/// Create a sequence file, inferring format and gzip compression from the path.
390///
391/// ```no_run
392/// let mut writer = fastx::create("contigs.fasta")?;
393/// writer.write_record(&fastx::Sequence::fasta("contig1", b"ACGT"))?;
394/// writer.finish()?;
395/// # Ok::<(), fastx::Error>(())
396/// ```
397pub fn create<P: AsRef<Path>>(path: P) -> Result<BoxedWriter> {
398    WriterBuilder::default().create(path)
399}
400
401/// Write records to standard output in `format`.
402pub fn stdout(format: Format) -> Result<BoxedWriter> {
403    WriterBuilder::default().format(format).stdout()
404}
405
406#[cfg(test)]
407mod tests {
408    use super::*;
409    use crate::reader::FastxReader;
410
411    fn write_to_string(format: Format, records: &[Sequence], width: usize) -> String {
412        let mut out = Vec::new();
413        let mut writer = FastxWriter::new(&mut out, format).line_width(width);
414        writer.write_all(records).unwrap();
415        writer.flush().unwrap();
416        String::from_utf8(out).unwrap()
417    }
418
419    #[test]
420    fn writes_wrapped_fasta() {
421        let records = vec![Sequence::fasta("a", b"ACGTACGTAC").with_description("d")];
422        assert_eq!(
423            write_to_string(Format::Fasta, &records, 4),
424            ">a d\nACGT\nACGT\nAC\n"
425        );
426        assert_eq!(
427            write_to_string(Format::Fasta, &records, 0),
428            ">a d\nACGTACGTAC\n"
429        );
430    }
431
432    #[test]
433    fn drops_quality_when_writing_fasta() {
434        let records = vec![Sequence::fastq("a", b"ACGT", b"IIII").unwrap()];
435        assert_eq!(write_to_string(Format::Fasta, &records, 60), ">a\nACGT\n");
436    }
437
438    #[test]
439    fn refuses_fastq_without_quality() {
440        let mut out = Vec::new();
441        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
442        let err = writer
443            .write_record(&Sequence::fasta("a", b"ACGT"))
444            .unwrap_err();
445        assert!(matches!(err, Error::MissingQuality { .. }));
446    }
447
448    #[test]
449    fn validates_on_demand() {
450        let mut out = Vec::new();
451        let mut writer = FastxWriter::new(&mut out, Format::Fasta).validate(Alphabet::Dna);
452        assert!(writer.write_record(&Sequence::fasta("a", b"ACGT")).is_ok());
453        let err = writer
454            .write_record(&Sequence::fasta("b", b"ACGX"))
455            .unwrap_err();
456        assert!(matches!(err, Error::InvalidByte { byte: b'X', .. }));
457        assert_eq!(writer.records_written(), 1);
458    }
459
460    #[test]
461    fn writes_parts_without_records() {
462        let mut out = Vec::new();
463        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
464        writer
465            .write_fastq("r1", Some("desc"), b"ACGT", b"IIII")
466            .unwrap();
467        writer.write_fasta("r2", None, b"TTTT").unwrap();
468        writer.flush().unwrap();
469        assert_eq!(out, b"@r1 desc\nACGT\n+\nIIII\n>r2\nTTTT\n");
470
471        let mut out = Vec::new();
472        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
473        assert!(writer.write_fastq("r", None, b"ACGT", b"II").is_err());
474    }
475
476    #[test]
477    fn round_trips_through_reader() {
478        let original: Vec<Sequence> = (0..50)
479            .map(|i| {
480                Sequence::fastq(
481                    format!("read{i}"),
482                    "ACGTN".repeat(i + 1).into_bytes(),
483                    "IIIII".repeat(i + 1).into_bytes(),
484                )
485                .unwrap()
486                .with_description(format!("record number {i}"))
487            })
488            .collect();
489
490        let text = write_to_string(Format::Fastq, &original, 60);
491        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
492            .collect::<Result<Vec<_>>>()
493            .unwrap();
494        assert_eq!(parsed, original);
495
496        // FASTA round trip re-joins wrapped lines.
497        let fasta_input: Vec<Sequence> =
498            original.iter().cloned().map(Sequence::into_fasta).collect();
499        let text = write_to_string(Format::Fasta, &fasta_input, 7);
500        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
501            .collect::<Result<Vec<_>>>()
502            .unwrap();
503        assert_eq!(parsed, fasta_input);
504    }
505
506    #[test]
507    fn empty_sequence_round_trips() {
508        let records = vec![
509            Sequence::fasta("a", Vec::new()),
510            Sequence::fasta("b", b"AC"),
511        ];
512        let text = write_to_string(Format::Fasta, &records, 60);
513        assert_eq!(text, ">a\n\n>b\nAC\n");
514        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
515            .collect::<Result<Vec<_>>>()
516            .unwrap();
517        assert_eq!(parsed, records);
518    }
519
520    #[test]
521    fn builder_requires_a_format() {
522        let mut out = Vec::new();
523        assert!(matches!(
524            WriterBuilder::default().build(&mut out),
525            Err(Error::UnknownFormat { .. })
526        ));
527        assert!(WriterBuilder::default()
528            .format(Format::Fasta)
529            .build(&mut out)
530            .is_ok());
531    }
532}
fastx/writer.rs

fastx/
writer.rs
