Propogate BrokenPipe error up to main and exit gracefully

3 files changed, 90 insertions, 52 deletions

diff --git a/src/choice.rs b/src/choice.rs
index 2bd0194..18f2a03 100644
--- a/src/choice.rs
+++ b/src/choice.rs
@@ -35,15 +35,20 @@ impl Choice {
         }
     }
 
-    pub fn print_choice<W: WriteReceiver>(&self, line: &str, config: &Config, handle: &mut W) {
+    pub fn print_choice<W: WriteReceiver>(
+        &self,
+        line: &str,
+        config: &Config,
+        handle: &mut W,
+    ) -> std::io::Result<()> {
         if config.opt.character_wise {
-            self.print_choice_generic(line.chars(), config, handle);
+            self.print_choice_generic(line.chars(), config, handle)
         } else {
             let line_iter = config
                 .separator
                 .split(line)
                 .filter(|s| !s.is_empty() || config.opt.non_greedy);
-            self.print_choice_generic(line_iter, config, handle);
+            self.print_choice_generic(line_iter, config, handle)
         }
     }
 
@@ -55,23 +60,30 @@ impl Choice {
         self.negative_index
     }
 
-    fn print_choice_generic<W, T, I>(&self, mut iter: I, config: &Config, handle: &mut W)
+    fn print_choice_generic<W, T, I>(
+        &self,
+        mut iter: I,
+        config: &Config,
+        handle: &mut W,
+    ) -> std::io::Result<()>
     where
         W: WriteReceiver,
         T: Writeable,
         I: Iterator<Item = T>,
     {
         if self.is_reverse_range() && !self.has_negative_index() {
-            self.print_choice_reverse(iter, config, handle);
+            self.print_choice_reverse(iter, config, handle)?;
         } else if self.has_negative_index() {
-            self.print_choice_negative(iter, config, handle);
+            self.print_choice_negative(iter, config, handle)?;
         } else {
             if self.start > 0 {
                 iter.nth((self.start - 1).try_into().unwrap());
             }
             let range = self.end.checked_sub(self.start).unwrap();
-            Choice::print_choice_loop_max_items(iter, config, handle, range);
+            Choice::print_choice_loop_max_items(iter, config, handle, range)?;
         }
+
+        Ok(())
     }
 
     fn print_choice_loop_max_items<W, T, I>(
@@ -79,7 +91,8 @@ impl Choice {
         config: &Config,
         handle: &mut W,
         max_items: isize,
-    ) where
+    ) -> std::io::Result<()>
+    where
         W: WriteReceiver,
         T: Writeable,
         I: Iterator<Item = T>,
@@ -88,14 +101,21 @@ impl Choice {
         for i in 0..=max_items {
             match peek_iter.next() {
                 Some(s) => {
-                    handle.write_choice(s, config, peek_iter.peek().is_some() && i != max_items);
+                    handle.write_choice(s, config, peek_iter.peek().is_some() && i != max_items)?;
                 }
                 None => break,
             };
         }
+
+        Ok(())
     }
 
-    fn print_choice_negative<W, T, I>(&self, iter: I, config: &Config, handle: &mut W)
+    fn print_choice_negative<W, T, I>(
+        &self,
+        iter: I,
+        config: &Config,
+        handle: &mut W,
+    ) -> std::io::Result<()>
     where
         W: WriteReceiver,
         T: Writeable,
@@ -106,21 +126,28 @@ impl Choice {
 
         if end > start {
             for word in vec[start..std::cmp::min(end, vec.len() - 1)].iter() {
-                handle.write_choice(*word, config, true);
+                handle.write_choice(*word, config, true)?;
             }
-            handle.write_choice(vec[std::cmp::min(end, vec.len() - 1)], config, false);
+            handle.write_choice(vec[std::cmp::min(end, vec.len() - 1)], config, false)?;
         } else if self.start < 0 {
             for word in vec[end + 1..=std::cmp::min(start, vec.len() - 1)]
                 .iter()
                 .rev()
             {
-                handle.write_choice(*word, config, true);
+                handle.write_choice(*word, config, true)?;
             }
-            handle.write_choice(vec[end], config, false);
+            handle.write_choice(vec[end], config, false)?;
         }
+
+        Ok(())
     }
 
-    fn print_choice_reverse<W, T, I>(&self, mut iter: I, config: &Config, handle: &mut W)
+    fn print_choice_reverse<W, T, I>(
+        &self,
+        mut iter: I,
+        config: &Config,
+        handle: &mut W,
+    ) -> std::io::Result<()>
     where
         W: WriteReceiver,
         T: Writeable,
@@ -145,10 +172,11 @@ impl Choice {
         let mut peek_iter = stack.iter().rev().peekable();
         loop {
             match peek_iter.next() {
-                Some(s) => handle.write_choice(*s, config, peek_iter.peek().is_some()),
+                Some(s) => handle.write_choice(*s, config, peek_iter.peek().is_some())?,
                 None => break,
             }
         }
+        Ok(())
     }
 
     fn get_negative_start_end<T>(&self, vec: &Vec<T>) -> (usize, usize) {
@@ -234,7 +262,9 @@ mod tests {
             let config = Config::from_iter(vec);
             let mut handle = BufWriter::new(MockStdout::new());
 
-            config.opt.choices[0].print_choice(&String::from(input), &config, &mut handle);
+            config.opt.choices[0]
+                .print_choice(&String::from(input), &config, &mut handle)
+                .unwrap();
 
             assert_eq!(
                 String::from(output),
@@ -258,22 +288,18 @@ mod tests {
             let mut handle = BufWriter::new(MockStdout::new());
             let mut handle1 = BufWriter::new(MockStdout::new());
 
-            config.opt.choices[0].print_choice(
-                &String::from("rust is pretty cool"),
-                &config,
-                &mut handle,
-            );
+            config.opt.choices[0]
+                .print_choice(&String::from("rust is pretty cool"), &config, &mut handle)
+                .unwrap();
 
             assert_eq!(
                 String::from("cool"),
                 MockStdout::str_from_buf_writer(handle)
             );
 
-            config.opt.choices[1].print_choice(
-                &String::from("rust is pretty cool"),
-                &config,
-                &mut handle1,
-            );
+            config.opt.choices[1]
+                .print_choice(&String::from("rust is pretty cool"), &config, &mut handle1)
+                .unwrap();
 
             assert_eq!(String::from("is"), MockStdout::str_from_buf_writer(handle1));
         }
@@ -747,9 +773,13 @@ mod tests {
         fn print_1_and_3_with_output_field_separator_rust_syntax_inclusive() {
             let config = Config::from_iter(vec!["choose", "1", "3", "-o", "#"]);
             let mut handle = BufWriter::new(MockStdout::new());
-            config.opt.choices[0].print_choice(&String::from("a b c d"), &config, &mut handle);
+            config.opt.choices[0]
+                .print_choice(&String::from("a b c d"), &config, &mut handle)
+                .unwrap();
             handle.write(&config.output_separator).unwrap();
-            config.opt.choices[1].print_choice(&String::from("a b c d"), &config, &mut handle);
+            config.opt.choices[1]
+                .print_choice(&String::from("a b c d"), &config, &mut handle)
+                .unwrap();
             assert_eq!(String::from("b#d"), MockStdout::str_from_buf_writer(handle));
         }
 
diff --git a/src/main.rs b/src/main.rs
index 5672889..1d00960 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,13 +25,25 @@ fn main() {
 
     let stdout = io::stdout();
     let lock = stdout.lock();
-    match opt.input {
+    let exit_result = match opt.input {
         Some(_) => main_generic(opt, &mut io::BufWriter::new(lock)),
         None => main_generic(opt, &mut io::LineWriter::new(lock)),
+    };
+
+    match exit_result {
+        Ok(_) => (),
+        Err(e) => {
+            if e.kind() == io::ErrorKind::BrokenPipe {
+                // BrokenPipe means whoever is reading the output hung up, we should
+                // gracefully exit
+            } else {
+                eprintln!("Failed to write to output: {}", e)
+            }
+        }
     }
 }
 
-fn main_generic<W: WriteReceiver>(opt: Opt, handle: &mut W) {
+fn main_generic<W: WriteReceiver>(opt: Opt, handle: &mut W) -> io::Result<()> {
     let config = Config::new(opt);
 
     let read = match &config.opt.input {
@@ -61,18 +73,17 @@ fn main_generic<W: WriteReceiver>(opt: Opt, handle: &mut W) {
                 let choice_iter = &mut config.opt.choices.iter().peekable();
 
                 while let Some(choice) = choice_iter.next() {
-                    choice.print_choice(l, &config, handle);
+                    choice.print_choice(l, &config, handle)?;
                     if choice_iter.peek().is_some() {
-                        handle.write_separator(&config);
+                        handle.write_separator(&config)?;
                     }
                 }
 
-                match handle.write(b"\n") {
-                    Ok(_) => (),
-                    Err(e) => eprintln!("Failed to write to output: {}", e),
-                }
+                handle.write(b"\n").map(|_| ())?
             }
             Err(e) => println!("Failed to read line: {}", e),
         }
     }
+
+    Ok(())
 }
diff --git a/src/writer.rs b/src/writer.rs
index cf30ee7..3b9e710 100644
--- a/src/writer.rs
+++ b/src/writer.rs
@@ -1,27 +1,24 @@
-use std::io::{BufWriter, LineWriter, Write};
+use std::io::{self, BufWriter, LineWriter, Write};
 
 use crate::config::Config;
 use crate::writeable::Writeable;
 
 pub trait WriteReceiver: Write {
-    fn write_choice<Wa: Writeable>(&mut self, b: Wa, config: &Config, print_separator: bool) {
-        let num_bytes_written = match self.write(&b.to_byte_buf()) {
-            Ok(x) => x,
-            Err(e) => {
-                eprintln!("Failed to write to output: {}", e);
-                0
-            }
-        };
+    fn write_choice<Wa: Writeable>(
+        &mut self,
+        b: Wa,
+        config: &Config,
+        print_separator: bool,
+    ) -> io::Result<()> {
+        let num_bytes_written = self.write(&b.to_byte_buf())?;
         if num_bytes_written > 0 && print_separator {
-            self.write_separator(config);
+            self.write_separator(config)?;
         };
+        Ok(())
     }
 
-    fn write_separator(&mut self, config: &Config) {
-        match self.write(&config.output_separator) {
-            Ok(_) => (),
-            Err(e) => eprintln!("Failed to write to output: {}", e),
-        }
+    fn write_separator(&mut self, config: &Config) -> io::Result<()> {
+        self.write(&config.output_separator).map(|_| ())
     }
 }