diff --git a/rust/rhg/src/ui.rs b/rust/rhg/src/ui.rs --- a/rust/rhg/src/ui.rs +++ b/rust/rhg/src/ui.rs @@ -1,7 +1,11 @@ use std::io; -use std::io::Write; +use std::io::{ErrorKind, Write}; -pub struct Ui {} +#[derive(Debug)] +pub struct Ui { + stdout: std::io::Stdout, + stderr: std::io::Stderr, +} /// The kind of user interface error pub enum UiError { @@ -14,20 +18,31 @@ /// The commandline user interface impl Ui { pub fn new() -> Self { - Ui {} + Ui { + stdout: std::io::stdout(), + stderr: std::io::stderr(), + } + } + + /// Returns a buffered handle on stdout for faster batch printing + /// operations. + pub fn stdout_buffer(&self) -> StdoutBuffer { + StdoutBuffer::new(self.stdout.lock()) } /// Write bytes to stdout pub fn write_stdout(&self, bytes: &[u8]) -> Result<(), UiError> { - let mut stdout = io::stdout(); + let mut stdout = self.stdout.lock(); self.write_stream(&mut stdout, bytes) - .or_else(|e| self.into_stdout_error(e))?; + .or_else(|e| self.handle_stdout_error(e))?; - stdout.flush().or_else(|e| self.into_stdout_error(e)) + stdout.flush().or_else(|e| self.handle_stdout_error(e)) } - fn into_stdout_error(&self, error: io::Error) -> Result<(), UiError> { + /// Sometimes writing to stdout is not possible, try writing to stderr to + /// signal that failure, otherwise just bail. + fn handle_stdout_error(&self, error: io::Error) -> Result<(), UiError> { self.write_stderr( &[b"abort: ", error.to_string().as_bytes(), b"\n"].concat(), )?; @@ -36,7 +51,7 @@ /// Write bytes to stderr pub fn write_stderr(&self, bytes: &[u8]) -> Result<(), UiError> { - let mut stderr = io::stderr(); + let mut stderr = self.stderr.lock(); self.write_stream(&mut stderr, bytes) .or_else(|e| Err(UiError::StderrError(e)))?; @@ -52,3 +67,43 @@ stream.write_all(bytes) } } + +/// A buffered stdout writer for faster batch printing operations. +pub struct StdoutBuffer { + buf: io::BufWriter, +} + +impl StdoutBuffer { + pub fn new(writer: W) -> Self { + let buf = io::BufWriter::new(writer); + Self { buf } + } + + /// Write bytes to stdout buffer + pub fn write_all(&mut self, bytes: &[u8]) -> Result<(), UiError> { + self.buf.write_all(bytes).or_else(|e| self.io_err(e)) + } + + /// Flush bytes to stdout + pub fn flush(&mut self) -> Result<(), UiError> { + self.buf.flush().or_else(|e| self.io_err(e)) + } + + fn io_err(&self, error: io::Error) -> Result<(), UiError> { + if let ErrorKind::BrokenPipe = error.kind() { + // This makes `| head` work for example + return Ok(()); + } + let mut stderr = io::stderr(); + + stderr + .write_all( + &[b"abort: ", error.to_string().as_bytes(), b"\n"].concat(), + ) + .map_err(|e| UiError::StderrError(e))?; + + stderr.flush().map_err(|e| UiError::StderrError(e))?; + + Err(UiError::StdoutError(error)) + } +}