From 8a66b6d65454c6eb156fdd29d9eb07a156de223c Mon Sep 17 00:00:00 2001
From: Sam Greenbury <sgreenbury@turing.ac.uk>
Date: Fri, 1 Nov 2024 09:49:16 +0000
Subject: [PATCH 1/2] Handle broken pipe

---
 popgetter_cli/src/display.rs |  3 ++-
 popgetter_cli/src/main.rs    | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/popgetter_cli/src/display.rs b/popgetter_cli/src/display.rs
index 85a6642..4df224e 100644
--- a/popgetter_cli/src/display.rs
+++ b/popgetter_cli/src/display.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::io::Write;
 use std::sync::OnceLock;
 
 use comfy_table::{presets::NOTHING, *};
@@ -192,7 +193,7 @@ pub fn display_search_results(
                 }
             }
         }
-        println!("\n{}", table);
+        writeln!(&mut std::io::stdout().lock(), "{}", table)?;
     }
     Ok(())
 }
diff --git a/popgetter_cli/src/main.rs b/popgetter_cli/src/main.rs
index 38df0b2..dedf05e 100644
--- a/popgetter_cli/src/main.rs
+++ b/popgetter_cli/src/main.rs
@@ -1,6 +1,8 @@
 mod cli;
 mod display;
 
+use std::io::prelude::*;
+
 use anyhow::Result;
 use clap::Parser;
 use cli::{Cli, RunCommand};
@@ -21,7 +23,16 @@ async fn main() -> Result<()> {
     debug!("config: {config:?}");
 
     if let Some(command) = args.command {
-        command.run(config).await?;
+        // Return ok if pipe is closed instead of error, otherwise return error
+        // See: https://stackoverflow.com/a/65760807, https://github.com/rust-lang/rust/issues/62569
+        if let Err(err) = command.run(config).await {
+            if let Some(err) = err.downcast_ref::<std::io::Error>() {
+                if err.kind() == std::io::ErrorKind::BrokenPipe {
+                    return Ok(());
+                }
+            }
+            Err(err)?;
+        }
     }
     Ok(())
 }

From 9717a6e19099de0d4576b649d956d807786fb99d Mon Sep 17 00:00:00 2001
From: Sam Greenbury <sgreenbury@turing.ac.uk>
Date: Fri, 1 Nov 2024 10:09:05 +0000
Subject: [PATCH 2/2] Refactor and replace println! with writeln!

---
 popgetter_cli/src/display.rs | 48 ++++++++++++------------------------
 1 file changed, 16 insertions(+), 32 deletions(-)

diff --git a/popgetter_cli/src/display.rs b/popgetter_cli/src/display.rs
index 4df224e..8b30a48 100644
--- a/popgetter_cli/src/display.rs
+++ b/popgetter_cli/src/display.rs
@@ -4,10 +4,7 @@ use std::sync::OnceLock;
 
 use comfy_table::{presets::NOTHING, *};
 use itertools::izip;
-use polars::{
-    frame::{DataFrame, UniqueKeepStrategy},
-    prelude::SortMultipleOptions,
-};
+use polars::{frame::DataFrame, prelude::SortMultipleOptions};
 use popgetter::{metadata::ExpandedMetadata, search::SearchResults, COL};
 
 static LOOKUP: OnceLock<HashMap<&'static str, &'static str>> = OnceLock::new();
@@ -109,8 +106,8 @@ pub fn display_countries(countries: DataFrame, max_results: Option<usize>) -> an
             country_iso3116_2.unwrap_or_default(),
         ]);
     }
-    println!("\n{}", table);
-    Ok(())
+
+    Ok(writeln!(&mut std::io::stdout(), "\n{}", table)?)
 }
 
 pub fn display_search_results(
@@ -193,7 +190,7 @@ pub fn display_search_results(
                 }
             }
         }
-        writeln!(&mut std::io::stdout().lock(), "{}", table)?;
+        writeln!(&mut std::io::stdout(), "{}", table)?;
     }
     Ok(())
 }
@@ -228,40 +225,29 @@ pub fn display_summary(results: SearchResults) -> anyhow::Result<()> {
     let column = table.column_mut(1).unwrap();
     column.set_cell_alignment(CellAlignment::Right);
 
-    println!("\n{}", table);
-    Ok(())
+    Ok(writeln!(&mut std::io::stdout(), "\n{}", table)?)
 }
 
 /// Display a given column from the search results
 pub fn display_column(search_results: SearchResults, column: &str) -> anyhow::Result<()> {
-    search_results
+    Ok(search_results
         .0
-        .select([column])?
+        .column(column)?
+        .rechunk()
         .iter()
-        .for_each(|series| {
-            series
-                .rechunk()
-                .iter()
-                .map(|el| el.get_str().map(|s| s.to_string()).unwrap())
-                .for_each(|el| println!("{el}"))
-        });
-    Ok(())
+        .map(|el| el.get_str().map(|s| s.to_string()).unwrap())
+        .try_for_each(|el| writeln!(&mut std::io::stdout(), "{el}"))?)
 }
 
 /// Display the unique values of a given column from the search results
 pub fn display_column_unique(search_results: SearchResults, column: &str) -> anyhow::Result<()> {
-    search_results
+    Ok(search_results
         .0
-        .select([column])?
-        .unique(None, UniqueKeepStrategy::Any, None)?
+        .column(column)?
+        .unique()?
         .iter()
-        .for_each(|series| {
-            series
-                .iter()
-                .map(|el| el.get_str().map(|s| s.to_string()).unwrap())
-                .for_each(|el| println!("{el}"))
-        });
-    Ok(())
+        .map(|el| el.get_str().map(|s| s.to_string()).unwrap())
+        .try_for_each(|el| writeln!(&mut std::io::stdout(), "{el}"))?)
 }
 
 /// Display the columns of the expanded metadata that can be used for displaying metrics results
@@ -272,7 +258,5 @@ pub fn display_metdata_columns(expanded_metadata: &ExpandedMetadata) -> anyhow::
         .collect()?
         .get_column_names()
         .into_iter()
-        .collect::<Vec<_>>()
-        .into_iter()
-        .for_each(|val| println!("{val}")))
+        .try_for_each(|el| writeln!(&mut std::io::stdout(), "{el}"))?)
 }