feat(history): Add new flag to allow custom output format (#662)

* feat(history): Add new flag to allow custom output format

* more efficient formatting

* add user and host

* docs

Co-authored-by: Conrad Ludgate <conrad.ludgate@truelayer.com>

3 files changed, 108 insertions, 31 deletions

diff --git a/src/command/client/history.rs b/src/command/client/history.rs
index 19a3df0e..145fe63a 100644
--- a/src/command/client/history.rs
+++ b/src/command/client/history.rs
@@ -1,11 +1,13 @@
 use std::{
     env,
+    fmt::{self, Display},
     io::{StdoutLock, Write},
     time::Duration,
 };
 
 use clap::Subcommand;
 use eyre::Result;
+use runtime_format::{FormatKey, FormatKeyError, ParsedFmt};
 
 use atuin_client::{
     database::{current_context, Database},
@@ -17,7 +19,7 @@ use atuin_client::{
 use atuin_client::sync;
 use log::debug;
 
-use super::search::format_duration;
+use super::search::format_duration_into;
 
 #[derive(Subcommand)]
 #[command(infer_subcommands = true)]
@@ -46,6 +48,11 @@ pub enum Cmd {
         /// Show only the text of the command
         #[arg(long)]
         cmd_only: bool,
+
+        /// Available variables: {command}, {directory}, {duration}, {user}, {host} and {time}.
+        /// Example: --format "{time} - [{duration}] - {directory}$\t{command}"
+        #[arg(long, short)]
+        format: Option<String>,
     },
 
     /// Get the last command ran
@@ -56,6 +63,11 @@ pub enum Cmd {
         /// Show only the text of the command
         #[arg(long)]
         cmd_only: bool,
+
+        /// Available variables: {command}, {directory}, {duration}, {user}, {host} and {time}.
+        /// Example: --format "{time} - [{duration}] - {directory}$\t{command}"
+        #[arg(long, short)]
+        format: Option<String>,
     },
 }
 
@@ -79,41 +91,74 @@ impl ListMode {
 }
 
 #[allow(clippy::cast_sign_loss)]
-pub fn print_list(h: &[History], list_mode: ListMode) {
+pub fn print_list(h: &[History], list_mode: ListMode, format: Option<&str>) {
     let w = std::io::stdout();
     let mut w = w.lock();
 
     match list_mode {
-        ListMode::Human => print_human_list(&mut w, h),
+        ListMode::Human => print_human_list(&mut w, h, format),
         ListMode::CmdOnly => print_cmd_only(&mut w, h),
-        ListMode::Regular => print_regular(&mut w, h),
+        ListMode::Regular => print_regular(&mut w, h, format),
     }
 
     w.flush().expect("failed to flush history");
 }
 
-#[allow(clippy::cast_sign_loss)]
-pub fn print_human_list(w: &mut StdoutLock, h: &[History]) {
-    for h in h.iter().rev() {
-        let duration = format_duration(Duration::from_nanos(std::cmp::max(h.duration, 0) as u64));
-
-        let time = h.timestamp.format("%Y-%m-%d %H:%M:%S");
-        let cmd = h.command.trim();
+/// type wrapper around `History` so we can implement traits
+struct FmtHistory<'a>(&'a History);
 
-        writeln!(w, "{time} · {duration}\t{cmd}").expect("failed to write history");
+/// defines how to format the history
+impl FormatKey for FmtHistory<'_> {
+    #[allow(clippy::cast_sign_loss)]
+    fn fmt(&self, key: &str, f: &mut fmt::Formatter<'_>) -> Result<(), FormatKeyError> {
+        match key {
+            "command" => f.write_str(self.0.command.trim())?,
+            "directory" => f.write_str(self.0.cwd.trim())?,
+            "duration" => {
+                let dur = Duration::from_nanos(std::cmp::max(self.0.duration, 0) as u64);
+                format_duration_into(dur, f)?;
+            }
+            "time" => self.0.timestamp.format("%Y-%m-%d %H:%M:%S").fmt(f)?,
+            "host" => f.write_str(
+                self.0
+                    .hostname
+                    .split_once(':')
+                    .map_or(&self.0.hostname, |(host, _)| host),
+            )?,
+            "user" => f.write_str(self.0.hostname.split_once(':').map_or("", |(_, user)| user))?,
+            _ => return Err(FormatKeyError::UnknownKey),
+        }
+        Ok(())
     }
 }
 
-#[allow(clippy::cast_sign_loss)]
-pub fn print_regular(w: &mut StdoutLock, h: &[History]) {
+fn print_list_with(w: &mut StdoutLock, h: &[History], format: &str) {
+    let fmt = match ParsedFmt::new(format) {
+        Ok(fmt) => fmt,
+        Err(err) => {
+            eprintln!("ERROR: History formatting failed with the following error: {err}");
+            println!("If your formatting string contains curly braces (eg: {{var}}) you need to escape them this way: {{{{var}}.");
+            std::process::exit(1)
+        }
+    };
+
     for h in h.iter().rev() {
-        let duration = format_duration(Duration::from_nanos(std::cmp::max(h.duration, 0) as u64));
+        writeln!(w, "{}", fmt.with_args(&FmtHistory(h))).expect("failed to write history");
+    }
+}
 
-        let time = h.timestamp.format("%Y-%m-%d %H:%M:%S");
-        let cmd = h.command.trim();
+pub fn print_human_list(w: &mut StdoutLock, h: &[History], format: Option<&str>) {
+    let format = format
+        .unwrap_or("{time} · {duration}\t{command}")
+        .replace("\\t", "\t");
+    print_list_with(w, h, &format);
+}
 
-        writeln!(w, "{time}\t{cmd}\t{duration}").expect("failed to write history");
-    }
+pub fn print_regular(w: &mut StdoutLock, h: &[History], format: Option<&str>) {
+    let format = format
+        .unwrap_or("{time}\t{command}\t{duration}")
+        .replace("\\t", "\t");
+    print_list_with(w, h, &format);
 }
 
 pub fn print_cmd_only(w: &mut StdoutLock, h: &[History]) {
@@ -187,6 +232,7 @@ impl Cmd {
                 cwd,
                 human,
                 cmd_only,
+                format,
             } => {
                 let session = if *session {
                     Some(env::var("ATUIN_SESSION")?)
@@ -218,14 +264,26 @@ impl Cmd {
                     }
                 };
 
-                print_list(&history, ListMode::from_flags(*human, *cmd_only));
+                print_list(
+                    &history,
+                    ListMode::from_flags(*human, *cmd_only),
+                    format.as_deref(),
+                );
 
                 Ok(())
             }
 
-            Self::Last { human, cmd_only } => {
+            Self::Last {
+                human,
+                cmd_only,
+                format,
+            } => {
                 let last = db.last().await?;
-                print_list(&[last], ListMode::from_flags(*human, *cmd_only));
+                print_list(
+                    &[last],
+                    ListMode::from_flags(*human, *cmd_only),
+                    format.as_deref(),
+                );
 
                 Ok(())
             }
diff --git a/src/command/client/search.rs b/src/command/client/search.rs
index e528576d..53471ec1 100644
--- a/src/command/client/search.rs
+++ b/src/command/client/search.rs
@@ -16,7 +16,7 @@ mod duration;
 mod event;
 mod history_list;
 mod interactive;
-pub use duration::format_duration;
+pub use duration::{format_duration, format_duration_into};
 
 #[allow(clippy::struct_excessive_bools)]
 #[derive(Parser)]
@@ -74,6 +74,11 @@ pub struct Cmd {
     /// Show only the text of the command
     #[arg(long)]
     cmd_only: bool,
+
+    /// Available variables: {command}, {directory}, {duration}, {user}, {host} and {time}.
+    /// Example: --format "{time} - [{duration}] - {directory}$\t{command}"
+    #[arg(long, short)]
+    format: Option<String>,
 }
 
 impl Cmd {
@@ -97,6 +102,7 @@ impl Cmd {
                 self.exit,
                 self.exclude_exit,
                 self.exclude_cwd,
+                self.format,
                 self.before,
                 self.after,
                 self.limit,
@@ -122,6 +128,7 @@ async fn run_non_interactive(
     exit: Option<i64>,
     exclude_exit: Option<i64>,
     exclude_cwd: Option<String>,
+    format: Option<String>,
     before: Option<String>,
     after: Option<String>,
     limit: Option<i64>,
@@ -202,6 +209,6 @@ async fn run_non_interactive(
         .map(std::borrow::ToOwned::to_owned)
         .collect();
 
-    super::history::print_list(&results, list_mode);
+    super::history::print_list(&results, list_mode, format.as_deref());
     Ok(results.len())
 }
diff --git a/src/command/client/search/duration.rs b/src/command/client/search/duration.rs
index 1dc4245f..08dadb95 100644
--- a/src/command/client/search/duration.rs
+++ b/src/command/client/search/duration.rs
@@ -1,10 +1,11 @@
+use core::fmt;
 use std::{ops::ControlFlow, time::Duration};
 
 #[allow(clippy::module_name_repetitions)]
-pub fn format_duration(f: Duration) -> String {
-    fn item(name: &str, value: u64) -> ControlFlow<String> {
+pub fn format_duration_into(dur: Duration, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    fn item(unit: &'static str, value: u64) -> ControlFlow<(&'static str, u64)> {
         if value > 0 {
-            ControlFlow::Break(format!("{value}{name}"))
+            ControlFlow::Break((unit, value))
         } else {
             ControlFlow::Continue(())
         }
@@ -13,7 +14,7 @@ pub fn format_duration(f: Duration) -> String {
     // impl taken and modified from
     // https://github.com/tailhook/humantime/blob/master/src/duration.rs#L295-L331
     // Copyright (c) 2016 The humantime Developers
-    fn fmt(f: Duration) -> ControlFlow<String, ()> {
+    fn fmt(f: Duration) -> ControlFlow<(&'static str, u64), ()> {
         let secs = f.as_secs();
         let nanos = f.subsec_nanos();
 
@@ -43,8 +44,19 @@ pub fn format_duration(f: Duration) -> String {
         ControlFlow::Continue(())
     }
 
-    match fmt(f) {
-        ControlFlow::Break(b) => b,
-        ControlFlow::Continue(()) => String::from("0s"),
+    match fmt(dur) {
+        ControlFlow::Break((unit, value)) => write!(f, "{value}{unit}"),
+        ControlFlow::Continue(()) => write!(f, "0s"),
+    }
+}
+
+#[allow(clippy::module_name_repetitions)]
+pub fn format_duration(f: Duration) -> String {
+    struct F(Duration);
+    impl fmt::Display for F {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            format_duration_into(self.0, f)
+        }
     }
+    F(f).to_string()
 }