Add WriterWithContext

willchet · willchet · commit c4acc6126901 · 2026-02-04T14:25:53.000-05:00
diff --git a/src/io/mod.rs b/src/io/mod.rs
@@ -1,4 +1,8 @@
-use std::{fmt::Display, io::Read, path::Path};
+use std::{
+    fmt::Display,
+    io::{Read, Write},
+    path::Path,
+};
 use zoe::{
     data::err::{ResultWithErrorContext, WithErrorContext},
     prelude::{FastQReader, FastaReader},
@@ -98,7 +102,7 @@ pub trait IterWithErrorContext: Sized {
     /// [`ErrorWithContext`]: zoe::data::err::ErrorWithContext
     fn iter_with_context(self, description: impl Into<String>) -> IterWithContext<Self>;
 
-    /// Convenience function for adding file context to an yielded errors.
+    /// Convenience function for adding file context to any yielded errors.
     ///
     /// The context will be formatted as `msg: file`. The `msg` field may be
     /// anything implementing [`Display`].
@@ -232,3 +236,73 @@ pub(crate) fn check_distinct_files(
 pub(crate) fn is_gz<P: AsRef<Path>>(path: P) -> bool {
     path.as_ref().extension().is_some_and(|ext| ext == "gz")
 }
+
+/// A wrapper around a writer of type `W` such that error context is added to
+/// any failed writes.
+#[derive(Debug)]
+pub struct WriterWithContext<W> {
+    /// The inner writer.
+    writer:      W,
+    /// The context to add to any failed writes.
+    description: String,
+}
+
+impl<W> Write for WriterWithContext<W>
+where
+    W: Write,
+{
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        Ok(self.writer.write(buf).with_context(&self.description)?)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(self.writer.flush().with_context(&self.description)?)
+    }
+
+    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
+        Ok(self.writer.write_vectored(bufs).with_context(&self.description)?)
+    }
+
+    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
+        Ok(self.writer.write_all(buf).with_context(&self.description)?)
+    }
+
+    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
+        Ok(self.writer.write_fmt(args).with_context(&self.description)?)
+    }
+}
+
+/// An extension trait for [`Write`] allowing additional context to be added to
+/// each failed write (via [`WriterWithContext`]).
+pub trait WriterWithErrorContext: Sized {
+    /// Wraps any errors that get produced during writing in an
+    /// [`ErrorWithContext`] with the given description.
+    ///
+    /// The `description` field may be anything implementing `Into<String>`.
+    /// Passing an owned `String` avoids an extra allocation.
+    ///
+    /// [`ErrorWithContext`]: zoe::data::err::ErrorWithContext
+    fn writer_with_context(self, description: impl Into<String>) -> WriterWithContext<Self>;
+
+    /// Convenience function for adding file context to any produced errors.
+    ///
+    /// The context will be formatted as `msg: file`. The `msg` field may be
+    /// anything implementing [`Display`].
+    fn writer_with_file_context(self, msg: impl Display, file: impl AsRef<Path>) -> WriterWithContext<Self>;
+}
+
+impl<W> WriterWithErrorContext for W
+where
+    W: Write,
+{
+    fn writer_with_context(self, description: impl Into<String>) -> WriterWithContext<Self> {
+        WriterWithContext {
+            writer:      self,
+            description: description.into(),
+        }
+    }
+
+    fn writer_with_file_context(self, msg: impl Display, file: impl AsRef<Path>) -> WriterWithContext<Self> {
+        Self::writer_with_context(self, format!("{msg}: '{path}'", path = file.as_ref().display()))
+    }
+}
diff --git a/src/io/open_options/context.rs b/src/io/open_options/context.rs
@@ -1,5 +1,5 @@
-use crate::io::{IterWithContext, IterWithErrorContext};
-use std::{error::Error, fmt::Display, path::Path};
+use crate::io::{IterWithContext, IterWithErrorContext, WriterWithContext, WriterWithErrorContext};
+use std::{error::Error, fmt::Display, io::Write, path::Path};
 use zoe::data::err::{ErrorWithContext, WithErrorContext};
 
 /// An enum to represent the possible reader types for [`InputOptions`], for the
@@ -203,11 +203,11 @@ impl InputContext<'_> {
 /// [`OutputOptions`]: crate::io::open_options::OutputOptions
 pub struct OutputContext<'a> {
     /// The [`OutputType`] used by the first output.
-    output1: OutputType<'a>,
+    pub output1: OutputType<'a>,
     /// The [`OutputType`] used by the second output. If there is no second
     /// output, this defaults to [`OutputType::Stdout`], but it will not
     /// influence the error context displayed.
-    output2: OutputType<'a>,
+    pub output2: OutputType<'a>,
 }
 
 impl<'a> OutputContext<'a> {
@@ -241,6 +241,17 @@ impl OutputContext<'_> {
             OutputType::Stdout => e.with_context("Failed to write to stdout"),
         }
     }
+
+    /// Given one of the `output` fields for an [`OutputContext`], add the
+    /// context to a provided writer.
+    pub fn add_writer_context<W>(writer: W, output: OutputType) -> WriterWithContext<W>
+    where
+        W: Write, {
+        match output {
+            OutputType::File(path) => writer.writer_with_file_context("Failed to write to file", path),
+            OutputType::Stdout => writer.writer_with_context("Failed to write to stdout"),
+        }
+    }
 }
 
 /// An error occurring while working with potentially-paired inputs or outputs,
diff --git a/src/io/open_options/output_options.rs b/src/io/open_options/output_options.rs
@@ -1,4 +1,7 @@
-use crate::io::{OptionalPaths, OutputContext, PairedErrors, RecordWriters, WriteFileZipStdout, open_options::PairedStruct};
+use crate::io::{
+    OptionalPaths, OutputContext, PairedErrors, RecordWriters, WriteFileZipStdout, WriterWithContext,
+    open_options::PairedStruct,
+};
 use std::{fs::File, io::BufWriter, path::Path};
 
 /// A builder pattern for creating output files in IRMA-core.
@@ -89,16 +92,8 @@ impl<'a> OutputOptions<'a, &'a Path> {
     }
 
     /// Interprets the path using [`File`] for writing.
-    #[allow(dead_code)]
-    pub fn use_file(self) -> OutputOptions<'a, BufWriter<File>> {
-        let output = self.output.and_then(|path| {
-            if let Some(capacity) = self.capacity {
-                File::create(path).map(|file| BufWriter::with_capacity(capacity, file))
-            } else {
-                File::create(path).map(BufWriter::new)
-            }
-            .map_err(PairedErrors::Err1)
-        });
+    pub fn use_file(self) -> OutputOptions<'a, File> {
+        let output = self.output.and_then(|path| File::create(path).map_err(PairedErrors::Err1));
 
         OutputOptions {
             context: self.context,
@@ -184,18 +179,10 @@ impl<'a> OutputOptions<'a, RecordWriters<&'a Path>> {
 
     /// Interprets the path(s) using [`File`] for writing.
     #[allow(dead_code)]
-    pub fn use_file(self) -> OutputOptions<'a, RecordWriters<BufWriter<File>>> {
+    pub fn use_file(self) -> OutputOptions<'a, RecordWriters<File>> {
         OutputOptions {
             context:  self.context,
-            output:   self.output.and_then(|writers| {
-                writers.try_map(|path| {
-                    if let Some(capacity) = self.capacity {
-                        File::create(path).map(|file| BufWriter::with_capacity(capacity, file))
-                    } else {
-                        File::create(path).map(BufWriter::new)
-                    }
-                })
-            }),
+            output:   self.output.and_then(|writers| writers.try_map(File::create)),
             capacity: self.capacity,
         }
     }
@@ -252,16 +239,26 @@ impl<'a> OutputOptions<'a, OptionalPaths<'a>> {
     }
 }
 
-impl<'a> OutputOptions<'a, BufWriter<File>> {
+impl<'a> OutputOptions<'a, File> {
     /// Opens the [`File`] for writing, wrapping it in a [`BufWriter`].
     ///
     /// ## Errors
     ///
     /// IO errors when opening the file are propagated. Context is added that
     /// includes the path.
-    pub fn open(self) -> std::io::Result<BufWriter<File>> {
-        match self.output {
-            Ok(writer) => Ok(writer),
+    pub fn open(self) -> std::io::Result<BufWriter<WriterWithContext<File>>> {
+        let output = self.output.map(|file| {
+            let output = OutputContext::add_writer_context(file, self.context.output1);
+
+            if let Some(capacity) = self.capacity {
+                BufWriter::with_capacity(capacity, output)
+            } else {
+                BufWriter::new(output)
+            }
+        });
+
+        match output {
+            Ok(output) => Ok(output),
             Err(e) => Err(self.context.add_context(e).into()),
         }
     }
@@ -282,7 +279,7 @@ impl<'a> OutputOptions<'a, WriteFileZipStdout> {
     }
 }
 
-impl<'a> OutputOptions<'a, RecordWriters<BufWriter<File>>> {
+impl<'a> OutputOptions<'a, RecordWriters<File>> {
     /// Opens the potentially paired [`File`]s for writing, wrapping each in a
     /// [`BufWriter`].
     ///
@@ -292,8 +289,18 @@ impl<'a> OutputOptions<'a, RecordWriters<BufWriter<File>>> {
     /// includes the path.
     #[allow(dead_code)]
     pub fn open(self) -> std::io::Result<RecordWriters<BufWriter<File>>> {
-        match self.output {
-            Ok(writer) => Ok(writer),
+        let output = self.output.map(|writers| {
+            writers.map(|file| {
+                if let Some(capacity) = self.capacity {
+                    BufWriter::with_capacity(capacity, file)
+                } else {
+                    BufWriter::new(file)
+                }
+            })
+        });
+
+        match output {
+            Ok(output) => Ok(output),
             Err(e) => Err(self.context.add_context(e).into()),
         }
     }
diff --git a/src/io/writers.rs b/src/io/writers.rs
@@ -1,4 +1,4 @@
-use crate::io::is_gz;
+use crate::io::{WriterWithContext, WriterWithErrorContext, is_gz};
 use flate2::{Compression, write::GzEncoder};
 use std::{
     fs::File,
@@ -9,15 +9,16 @@ use zoe::define_whichever;
 
 define_whichever! {
     /// An enum for the different acceptable output types. A [`BufWriter`] is
-    /// used for all variants.
+    /// used for all variants, and all variants are wrapped in
+    /// [`WriterWithContext`] to context to write errors.
     #[derive(Debug)]
     pub(crate) enum  WriteFileZipStdout {
         /// A writer for a regular uncompressed file.
-        File(BufWriter<File>),
+        File(WriterWithContext<BufWriter<File>>),
         /// A writer for a gzip compressed file.
-        Zipped(GzEncoder<BufWriter<File>>),
+        Zipped(WriterWithContext<GzEncoder<BufWriter<File>>>),
         /// A writer for uncompressed data to stdout.
-        Stdout(BufWriter<Stdout>),
+        Stdout(WriterWithContext<BufWriter<Stdout>>),
     }
 
     impl Write for WriteFileZipStdout {}
@@ -37,15 +38,20 @@ impl WriteFileZipStdout {
                 let file = File::create(&path)?;
                 let bufwriter = BufWriter::new(file);
 
-                let writer = if is_gz(path) {
-                    Self::Zipped(GzEncoder::new(bufwriter, Compression::default()))
+                let writer = if is_gz(&path) {
+                    Self::Zipped(
+                        GzEncoder::new(bufwriter, Compression::default())
+                            .writer_with_file_context("Failed to write to zipped file", path),
+                    )
                 } else {
-                    Self::File(bufwriter)
+                    Self::File(bufwriter.writer_with_file_context("Failed to write to file", path))
                 };
 
                 Ok(writer)
             }
-            None => Ok(WriteFileZipStdout::Stdout(BufWriter::new(stdout()))),
+            None => Ok(WriteFileZipStdout::Stdout(
+                BufWriter::new(stdout()).writer_with_context("Failed to write to stdout"),
+            )),
         }
     }
 
@@ -57,15 +63,20 @@ impl WriteFileZipStdout {
                 let file = File::create(&path)?;
                 let bufwriter = BufWriter::with_capacity(capacity, file);
 
-                let writer = if is_gz(path) {
-                    Self::Zipped(GzEncoder::new(bufwriter, Compression::default()))
+                let writer = if is_gz(&path) {
+                    Self::Zipped(
+                        GzEncoder::new(bufwriter, Compression::default())
+                            .writer_with_file_context("Failed to write to zipped file", path),
+                    )
                 } else {
-                    Self::File(bufwriter)
+                    Self::File(bufwriter.writer_with_file_context("Failed to write to file", path))
                 };
 
                 Ok(writer)
             }
-            None => Ok(WriteFileZipStdout::Stdout(BufWriter::with_capacity(capacity, stdout()))),
+            None => Ok(WriteFileZipStdout::Stdout(
+                BufWriter::with_capacity(capacity, stdout()).writer_with_context("Failed to write to stdout"),
+            )),
         }
     }
 }
diff --git a/src/processes/integrated/preprocess.rs b/src/processes/integrated/preprocess.rs
@@ -4,7 +4,7 @@
 
 use crate::{
     args::clipping::{ClippingArgs, ParsedClippingArgs, parse_clipping_args},
-    io::{InputOptions, IterWithContext, OutputOptions, ReadFileZipPipe, RecordReaders},
+    io::{InputOptions, IterWithContext, OutputOptions, ReadFileZipPipe, RecordReaders, WriterWithContext},
     qc::{fastq::ReadTransforms, fastq_metadata::*},
     utils::{
         get_hasher,
@@ -103,10 +103,10 @@ pub fn preprocess_process(args: PreprocessArgs) -> Result<(), std::io::Error> {
 }
 
 struct ParsedPreprocessIoArgs {
-    table_writer: BufWriter<File>,
+    table_writer: BufWriter<WriterWithContext<File>>,
     reader1:      IterWithContext<FastQReader<ReadFileZipPipe>>,
     reader2:      Option<IterWithContext<FastQReader<ReadFileZipPipe>>>,
-    log_writer:   Option<BufWriter<File>>,
+    log_writer:   Option<BufWriter<WriterWithContext<File>>>,
     log_file:     Option<PathBuf>,
 }
 
@@ -253,7 +253,7 @@ fn trim_and_deflate(
 /// Writes the table file to `table_writer` and the XFL file to STDOUT. The
 /// number of read patterns is returned.
 fn output_deflated_sequences(
-    metadata_by_sequence: DeflatedSequences, mut table_writer: BufWriter<File>,
+    metadata_by_sequence: DeflatedSequences, mut table_writer: impl Write,
 ) -> Result<usize, std::io::Error> {
     let mut stdout_writer = BufWriter::new(std::io::stdout());
 
@@ -282,7 +282,7 @@ fn output_deflated_sequences(
 
 /// Writes the log file.
 fn write_log(
-    mut log_writer: BufWriter<File>, metadata: &FastQMetadata, paired_reads: bool, read_pattern_count_passing: usize,
+    mut log_writer: impl Write, metadata: &FastQMetadata, paired_reads: bool, read_pattern_count_passing: usize,
     options: &ParsedPreprocessOptions, log_file: PathBuf,
 ) -> Result<(), std::io::Error> {
     let FastQMetadata {
