Use buffered writer in args::parse_arguments

To be able to decide whether to print the argparse output depending on
the result of the argument parsing, this patch wraps stdout and stderr
in a BufWriter before invoking argparse. Our BufWriter implementation
only writes to the inner Write if the flush method is called. This
allows us to decide whether the buffered data should be written or
silently dropped.

1 files changed, 52 insertions, 1 deletions

diff --git a/nitrocli/src/args.rs b/nitrocli/src/args.rs
index a0ad583..924b10e 100644
--- a/nitrocli/src/args.rs
+++ b/nitrocli/src/args.rs
@@ -29,6 +29,38 @@ use crate::RunCtx;
 
 type Result<T> = result::Result<T, Error>;
 
+/// Wraps a writer and buffers its output.
+///
+/// This implementation is similar to `io::BufWriter`, but:
+/// - The inner writer is only written to if `flush` is called.
+/// - The buffer may grow infinitely large.
+struct BufWriter<'w, W: io::Write + ?Sized> {
+  buf: Vec<u8>,
+  inner: &'w mut W,
+}
+
+impl<'w, W: io::Write + ?Sized> BufWriter<'w, W> {
+  pub fn new(inner: &'w mut W) -> Self {
+    BufWriter {
+      buf: Vec::with_capacity(128),
+      inner,
+    }
+  }
+}
+
+impl<'w, W: io::Write + ?Sized> io::Write for BufWriter<'w, W> {
+  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+    self.buf.extend_from_slice(buf);
+    Ok(buf.len())
+  }
+
+  fn flush(&mut self) -> io::Result<()> {
+    self.inner.write_all(&self.buf)?;
+    self.buf.clear();
+    self.inner.flush()
+  }
+}
+
 trait Stdio {
   fn stdio(&mut self) -> (&mut dyn io::Write, &mut dyn io::Write);
 }
@@ -39,6 +71,15 @@ impl<'io> Stdio for RunCtx<'io> {
   }
 }
 
+impl<W> Stdio for (&mut W, &mut W)
+where
+  W: io::Write,
+{
+  fn stdio(&mut self) -> (&mut dyn io::Write, &mut dyn io::Write) {
+    (self.0, self.1)
+  }
+}
+
 /// A command execution context that captures additional data pertaining
 /// the command execution.
 pub struct ExecCtx<'io> {
@@ -827,6 +868,8 @@ fn parse_arguments<'io, 'ctx: 'io>(
   ctx: &'ctx mut RunCtx<'_>,
   args: Vec<String>,
 ) -> Result<(Command, ExecCtx<'io>, Vec<String>)> {
+  use std::io::Write;
+
   let mut model: Option<DeviceModel> = None;
   let model_help = format!(
     "Select the device model to connect to ({})",
@@ -862,8 +905,16 @@ fn parse_arguments<'io, 'ctx: 'io>(
     "The arguments for the command",
   );
   parser.stop_on_first_argument(true);
-  parse(ctx, parser, args)?;
 
+  let mut stdout_buf = BufWriter::new(ctx.stdout);
+  let mut stderr_buf = BufWriter::new(ctx.stderr);
+  let mut stdio_buf = (&mut stdout_buf, &mut stderr_buf);
+  let result = parse(&mut stdio_buf, parser, args);
+
+  stdout_buf.flush()?;
+  stderr_buf.flush()?;
+
+  result?;
   subargs.insert(0, format!("nitrocli {}", command));
 
   let ctx = ExecCtx {