From f8d01eef998c86c52514896539c13dbfe1837e55 Mon Sep 17 00:00:00 2001
From: David <drmorr@appliedcomputing.io>
Date: Thu, 15 Feb 2024 10:52:19 -0800
Subject: [PATCH] feat: add 'ignored_commands' option to stats (#1722)

---
 atuin-client/config.toml          |  3 ++
 atuin-client/src/settings.rs      |  7 +++++
 atuin/src/command/client/stats.rs | 46 +++++++++++++++++++++++++++----
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/atuin-client/config.toml b/atuin-client/config.toml
index facc0bd4..19a6d61e 100644
--- a/atuin-client/config.toml
+++ b/atuin-client/config.toml
@@ -171,3 +171,6 @@ enter_accept = true
 
 ## Set commands that should be totally stripped and ignored from stats
 #common_prefix = ["sudo"]
+
+## Set commands that will be completely ignored from stats
+#ignored_commands = ["cd", "ls", "vi"]
diff --git a/atuin-client/src/settings.rs b/atuin-client/src/settings.rs
index f4c47c64..763fde1b 100644
--- a/atuin-client/src/settings.rs
+++ b/atuin-client/src/settings.rs
@@ -270,6 +270,8 @@ pub struct Stats {
     pub common_prefix: Vec<String>, // sudo, etc. commands we want to strip off
     #[serde(default = "Stats::common_subcommands_default")]
     pub common_subcommands: Vec<String>, // kubectl, commands we should consider subcommands for
+    #[serde(default = "Stats::ignored_commands_default")]
+    pub ignored_commands: Vec<String>, // cd, ls, etc. commands we want to completely hide from stats
 }
 
 impl Stats {
@@ -283,6 +285,10 @@ impl Stats {
             .map(String::from)
             .collect()
     }
+
+    fn ignored_commands_default() -> Vec<String> {
+        vec![]
+    }
 }
 
 impl Default for Stats {
@@ -290,6 +296,7 @@ impl Default for Stats {
         Self {
             common_prefix: Self::common_prefix_default(),
             common_subcommands: Self::common_subcommands_default(),
+            ignored_commands: Self::ignored_commands_default(),
         }
     }
 }
diff --git a/atuin/src/command/client/stats.rs b/atuin/src/command/client/stats.rs
index f73b8bbb..c2636536 100644
--- a/atuin/src/command/client/stats.rs
+++ b/atuin/src/command/client/stats.rs
@@ -24,16 +24,22 @@ pub struct Cmd {
     count: usize,
 }
 
-fn compute_stats(settings: &Settings, history: &[History], count: usize) {
+fn compute_stats(settings: &Settings, history: &[History], count: usize) -> (usize, usize) {
     let mut commands = HashSet::<&str>::with_capacity(history.len());
     let mut prefixes = HashMap::<&str, usize>::with_capacity(history.len());
+    let mut total_unignored = 0;
     for i in history {
         // just in case it somehow has a leading tab or space or something (legacy atuin didn't ignore space prefixes)
         let command = i.command.trim();
+        let prefix = interesting_command(settings, command);
+
+        if settings.stats.ignored_commands.iter().any(|c| c == prefix) {
+            continue;
+        }
+
+        total_unignored += 1;
         commands.insert(command);
-        *prefixes
-            .entry(interesting_command(settings, command))
-            .or_default() += 1;
+        *prefixes.entry(prefix).or_default() += 1;
     }
 
     let unique = commands.len();
@@ -42,7 +48,7 @@ fn compute_stats(settings: &Settings, history: &[History], count: usize) {
     top.truncate(count);
     if top.is_empty() {
         println!("No commands found");
-        return;
+        return (0, 0);
     }
 
     let max = top.iter().map(|x| x.1).max().unwrap();
@@ -73,8 +79,10 @@ fn compute_stats(settings: &Settings, history: &[History], count: usize) {
             command.escape_control()
         );
     }
-    println!("Total commands:   {}", history.len());
+    println!("Total commands:   {total_unignored}");
     println!("Unique commands:  {unique}");
+
+    (total_unignored, unique)
 }
 
 impl Cmd {
@@ -176,10 +184,36 @@ fn interesting_command<'a>(settings: &Settings, mut command: &'a str) -> &'a str
 
 #[cfg(test)]
 mod tests {
+    use atuin_client::history::History;
     use atuin_client::settings::Settings;
+    use time::OffsetDateTime;
 
+    use super::compute_stats;
     use super::interesting_command;
 
+    #[test]
+    fn ignored_commands() {
+        let mut settings = Settings::utc();
+        settings.stats.ignored_commands.push("cd".to_string());
+
+        let history = [
+            History::import()
+                .timestamp(OffsetDateTime::now_utc())
+                .command("cd foo")
+                .build()
+                .into(),
+            History::import()
+                .timestamp(OffsetDateTime::now_utc())
+                .command("cargo build stuff")
+                .build()
+                .into(),
+        ];
+
+        let (total, unique) = compute_stats(&settings, &history, 10);
+        assert_eq!(total, 1);
+        assert_eq!(unique, 1);
+    }
+
     #[test]
     fn interesting_commands() {
         let settings = Settings::utc();