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            Compression::Zstd => zstd_writer(buffered, self.level)?,
318        };
319        Ok(FastxWriter {
320            inner: sink,
321            format,
322            line_width: self.line_width,
323            validate: self.validate,
324            written: 0,
325        })
326    }
327
328    /// Write to standard output.
329    pub fn stdout(&self) -> Result<BoxedWriter> {
330        let format = self.format.unwrap_or(Format::Fasta);
331        let sink: Box<dyn Write + Send> = match self.compression.unwrap_or(Compression::None) {
332            Compression::None => Box::new(BufWriter::with_capacity(self.buffer_size, io::stdout())),
333            Compression::Gzip => gzip_writer(
334                BufWriter::with_capacity(self.buffer_size, io::stdout()),
335                self.level,
336            )?,
337            Compression::Bgzf => bgzf_writer(
338                BufWriter::with_capacity(self.buffer_size, io::stdout()),
339                self.level,
340            )?,
341            Compression::Zstd => zstd_writer(
342                BufWriter::with_capacity(self.buffer_size, io::stdout()),
343                self.level,
344            )?,
345        };
346        Ok(FastxWriter {
347            inner: sink,
348            format,
349            line_width: self.line_width,
350            validate: self.validate,
351            written: 0,
352        })
353    }
354}
355
356#[cfg(feature = "gzip")]
357fn gzip_writer<W: Write + Send + 'static>(
358    sink: W,
359    level: CompressionLevel,
360) -> Result<Box<dyn Write + Send>> {
361    Ok(Box::new(flate2::write::GzEncoder::new(
362        sink,
363        flate2::Compression::new(level.0.min(9)),
364    )))
365}
366
367#[cfg(not(feature = "gzip"))]
368fn gzip_writer<W: Write + Send + 'static>(
369    _sink: W,
370    _level: CompressionLevel,
371) -> Result<Box<dyn Write + Send>> {
372    Err(Error::FeatureDisabled("gzip"))
373}
374
375#[cfg(feature = "gzip")]
376fn bgzf_writer<W: Write + Send + 'static>(
377    sink: W,
378    level: CompressionLevel,
379) -> Result<Box<dyn Write + Send>> {
380    Ok(Box::new(crate::bgzf::BgzfWriter::with_level(sink, level)))
381}
382
383#[cfg(not(feature = "gzip"))]
384fn bgzf_writer<W: Write + Send + 'static>(
385    _sink: W,
386    _level: CompressionLevel,
387) -> Result<Box<dyn Write + Send>> {
388    Err(Error::FeatureDisabled("gzip"))
389}
390
391/// Zstandard output.
392///
393/// `auto_finish` matters: a zstd frame needs an epilogue, and without it the
394/// file is truncated. This mirrors how flate2's encoder finishes on drop, so
395/// both compressed paths behave the same whether the caller calls `finish` or
396/// simply drops the writer.
397#[cfg(feature = "zstd")]
398fn zstd_writer<W: Write + Send + 'static>(
399    sink: W,
400    level: CompressionLevel,
401) -> Result<Box<dyn Write + Send>> {
402    // gzip levels run 0-9, zstd 1-22. Map the level across rather than passing
403    // a gzip number to zstd, where 9 would be a middling setting.
404    let zstd_level = match level.0 {
405        0 => 1,
406        level => (level.min(9) as i32 - 1) * 21 / 8 + 1,
407    };
408    let encoder = zstd::stream::write::Encoder::new(sink, zstd_level)?;
409    Ok(Box::new(encoder.auto_finish()))
410}
411
412#[cfg(not(feature = "zstd"))]
413fn zstd_writer<W: Write + Send + 'static>(
414    _sink: W,
415    _level: CompressionLevel,
416) -> Result<Box<dyn Write + Send>> {
417    Err(Error::FeatureDisabled("zstd"))
418}
419
420/// The boxed writer type produced by [`create`] and [`WriterBuilder::stdout`].
421pub type BoxedWriter = FastxWriter<Box<dyn Write + Send>>;
422
423/// Create a sequence file, inferring format and gzip compression from the path.
424///
425/// ```no_run
426/// let mut writer = fastx::create("contigs.fasta")?;
427/// writer.write_record(&fastx::Sequence::fasta("contig1", b"ACGT"))?;
428/// writer.finish()?;
429/// # Ok::<(), fastx::Error>(())
430/// ```
431pub fn create<P: AsRef<Path>>(path: P) -> Result<BoxedWriter> {
432    WriterBuilder::default().create(path)
433}
434
435/// Write records to standard output in `format`.
436pub fn stdout(format: Format) -> Result<BoxedWriter> {
437    WriterBuilder::default().format(format).stdout()
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use crate::reader::FastxReader;
444
445    fn write_to_string(format: Format, records: &[Sequence], width: usize) -> String {
446        let mut out = Vec::new();
447        let mut writer = FastxWriter::new(&mut out, format).line_width(width);
448        writer.write_all(records).unwrap();
449        writer.flush().unwrap();
450        String::from_utf8(out).unwrap()
451    }
452
453    #[test]
454    fn writes_wrapped_fasta() {
455        let records = vec![Sequence::fasta("a", b"ACGTACGTAC").with_description("d")];
456        assert_eq!(
457            write_to_string(Format::Fasta, &records, 4),
458            ">a d\nACGT\nACGT\nAC\n"
459        );
460        assert_eq!(
461            write_to_string(Format::Fasta, &records, 0),
462            ">a d\nACGTACGTAC\n"
463        );
464    }
465
466    #[test]
467    fn drops_quality_when_writing_fasta() {
468        let records = vec![Sequence::fastq("a", b"ACGT", b"IIII").unwrap()];
469        assert_eq!(write_to_string(Format::Fasta, &records, 60), ">a\nACGT\n");
470    }
471
472    #[test]
473    fn refuses_fastq_without_quality() {
474        let mut out = Vec::new();
475        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
476        let err = writer
477            .write_record(&Sequence::fasta("a", b"ACGT"))
478            .unwrap_err();
479        assert!(matches!(err, Error::MissingQuality { .. }));
480    }
481
482    #[test]
483    fn validates_on_demand() {
484        let mut out = Vec::new();
485        let mut writer = FastxWriter::new(&mut out, Format::Fasta).validate(Alphabet::Dna);
486        assert!(writer.write_record(&Sequence::fasta("a", b"ACGT")).is_ok());
487        let err = writer
488            .write_record(&Sequence::fasta("b", b"ACGX"))
489            .unwrap_err();
490        assert!(matches!(err, Error::InvalidByte { byte: b'X', .. }));
491        assert_eq!(writer.records_written(), 1);
492    }
493
494    #[test]
495    fn writes_parts_without_records() {
496        let mut out = Vec::new();
497        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
498        writer
499            .write_fastq("r1", Some("desc"), b"ACGT", b"IIII")
500            .unwrap();
501        writer.write_fasta("r2", None, b"TTTT").unwrap();
502        writer.flush().unwrap();
503        assert_eq!(out, b"@r1 desc\nACGT\n+\nIIII\n>r2\nTTTT\n");
504
505        let mut out = Vec::new();
506        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
507        assert!(writer.write_fastq("r", None, b"ACGT", b"II").is_err());
508    }
509
510    #[test]
511    fn round_trips_through_reader() {
512        let original: Vec<Sequence> = (0..50)
513            .map(|i| {
514                Sequence::fastq(
515                    format!("read{i}"),
516                    "ACGTN".repeat(i + 1).into_bytes(),
517                    "IIIII".repeat(i + 1).into_bytes(),
518                )
519                .unwrap()
520                .with_description(format!("record number {i}"))
521            })
522            .collect();
523
524        let text = write_to_string(Format::Fastq, &original, 60);
525        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
526            .collect::<Result<Vec<_>>>()
527            .unwrap();
528        assert_eq!(parsed, original);
529
530        // FASTA round trip re-joins wrapped lines.
531        let fasta_input: Vec<Sequence> =
532            original.iter().cloned().map(Sequence::into_fasta).collect();
533        let text = write_to_string(Format::Fasta, &fasta_input, 7);
534        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
535            .collect::<Result<Vec<_>>>()
536            .unwrap();
537        assert_eq!(parsed, fasta_input);
538    }
539
540    #[test]
541    fn empty_sequence_round_trips() {
542        let records = vec![
543            Sequence::fasta("a", Vec::new()),
544            Sequence::fasta("b", b"AC"),
545        ];
546        let text = write_to_string(Format::Fasta, &records, 60);
547        assert_eq!(text, ">a\n\n>b\nAC\n");
548        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
549            .collect::<Result<Vec<_>>>()
550            .unwrap();
551        assert_eq!(parsed, records);
552    }
553
554    #[test]
555    fn builder_requires_a_format() {
556        let mut out = Vec::new();
557        assert!(matches!(
558            WriterBuilder::default().build(&mut out),
559            Err(Error::UnknownFormat { .. })
560        ));
561        assert!(WriterBuilder::default()
562            .format(Format::Fasta)
563            .build(&mut out)
564            .is_ok());
565    }
566}
fastx/writer.rs

fastx/
writer.rs
