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 {