diff --git a/src/lib.rs b/src/lib.rs index 5fe93e3..ee03df7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,36 @@ fn to_c_char(text: &str) -> *mut c_char { return CString::new(text).unwrap().into_raw(); } +enum Output { + Stdout(String), + Stderr(String), +} + +struct OutputCapture { + stdout: BufferRedirect, + stderr: BufferRedirect, +} + +impl OutputCapture { + pub fn capture() -> std::io::Result { + Ok(Self { + stdout: BufferRedirect::stdout()?, + stderr: BufferRedirect::stderr()?, + }) + } + + pub fn get_output(mut self) -> std::io::Result { + let mut buf = String::new(); + + if self.stderr.read_to_string(&mut buf)? > 0 { + Ok(Output::Stderr(buf)) + } else { + self.stdout.read_to_string(&mut buf)?; + Ok(Output::Stdout(buf)) + } + } +} + #[derive(Debug)] pub struct MenuError { message: String, @@ -154,19 +184,15 @@ impl Menu { } pub fn run(&self) -> Result { + let output = OutputCapture::capture().unwrap(); + unsafe { - let mut stdout = BufferRedirect::stdout().unwrap(); - match bindings::menu_run(self.ctx) { - 0 => { - let mut buf = String::new(); - stdout.read_to_string(&mut buf).unwrap(); - Ok(buf.trim_end().to_string()) - } - n => { - let mut buf = String::new(); - stdout.read_to_string(&mut buf).unwrap(); - Err(MenuError::new(n, buf)) + match (bindings::menu_run(self.ctx), output.get_output()) { + (0, Ok(Output::Stdout(value))) => Ok(value.trim_end().to_string()), + (n, Ok(Output::Stdout(value) | Output::Stderr(value))) => { + Err(MenuError::new(n, value)) } + (n, Err(e)) => Err(MenuError::new(n, e.to_string())), } } } @@ -197,6 +223,7 @@ mod tests { .add_items(&items); let result = menu.run().expect("couldn't get result"); + println!("{:?}", result.as_bytes()); assert!(items.contains(&result.as_str())); } }