From 0f207cd70126a717e8a8d301daca55b4c0a9997f Mon Sep 17 00:00:00 2001
From: rowan <rowan@kitsu.cafe>
Date: Sun, 26 Jan 2025 23:08:57 -0600
Subject: [PATCH] chore: refactor stdout capture

---
 Cargo.lock |  3 +--
 src/lib.rs | 49 ++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b496767..88cc4cd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -635,11 +635,10 @@ dependencies = [
 ]
 
 [[package]]
-name = "wmenu"
+name = "wmenu-rs"
 version = "0.1.0"
 dependencies = [
  "gag",
- "pangocairo-sys",
  "wmenu-sys",
 ]
 
diff --git a/src/lib.rs b/src/lib.rs
index 1d36036..36e003a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -20,6 +20,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<Self> {
+        Ok(Self {
+            stdout: BufferRedirect::stdout()?,
+            stderr: BufferRedirect::stderr()?,
+        })
+    }
+
+    pub fn get_output(mut self) -> std::io::Result<Output> {
+        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,
@@ -152,19 +182,15 @@ impl Menu {
     }
 
     pub fn run(&self) -> Result<String, MenuError> {
+        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())),
             }
         }
     }
@@ -195,6 +221,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()));
     }
 }