feat: allow configuring stats prefix (#1411)

2 files changed, 50 insertions, 14 deletions

diff --git a/atuin-client/src/settings.rs b/atuin-client/src/settings.rs
index fca76004..c5e72025 100644
--- a/atuin-client/src/settings.rs
+++ b/atuin-client/src/settings.rs
@@ -144,6 +144,24 @@ pub enum WordJumpMode {
 }
 
 #[derive(Clone, Debug, Deserialize)]
+pub struct Stats {
+    pub common_prefix: Vec<String>, // sudo, etc. commands we want to strip off
+    pub common_subcommands: Vec<String>, // kubectl, commands we should consider subcommands for
+}
+
+impl Default for Stats {
+    fn default() -> Self {
+        Self {
+            common_prefix: vec!["sudo", "doas"].into_iter().map(String::from).collect(),
+            common_subcommands: vec!["cargo", "go", "git", "npm", "yarn", "pnpm", "kubectl"]
+                .into_iter()
+                .map(String::from)
+                .collect(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
 pub struct Settings {
     pub dialect: Dialect,
     pub style: Style,
@@ -169,10 +187,13 @@ pub struct Settings {
     pub word_jump_mode: WordJumpMode,
     pub word_chars: String,
     pub scroll_context_lines: usize,
+
     #[serde(with = "serde_regex", default = "RegexSet::empty")]
     pub history_filter: RegexSet,
+
     #[serde(with = "serde_regex", default = "RegexSet::empty")]
     pub cwd_filter: RegexSet,
+
     pub secrets_filter: bool,
     pub workspaces: bool,
     pub ctrl_n_shortcuts: bool,
@@ -181,6 +202,9 @@ pub struct Settings {
     pub network_timeout: u64,
     pub enter_accept: bool,
 
+    #[serde(default)]
+    pub stats: Stats,
+
     // This is automatically loaded when settings is created. Do not set in
     // config! Keep secrets and settings apart.
     pub session_token: String,
diff --git a/atuin/src/command/client/stats.rs b/atuin/src/command/client/stats.rs
index 625c6beb..27ff362d 100644
--- a/atuin/src/command/client/stats.rs
+++ b/atuin/src/command/client/stats.rs
@@ -23,14 +23,16 @@ pub struct Cmd {
     count: usize,
 }
 
-fn compute_stats(history: &[History], count: usize) -> Result<()> {
+fn compute_stats(settings: &Settings, history: &[History], count: usize) -> Result<()> {
     let mut commands = HashSet::<&str>::with_capacity(history.len());
     let mut prefixes = HashMap::<&str, usize>::with_capacity(history.len());
     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();
         commands.insert(command);
-        *prefixes.entry(interesting_command(command)).or_default() += 1;
+        *prefixes
+            .entry(interesting_command(settings, command))
+            .or_default() += 1;
     }
 
     let unique = commands.len();
@@ -109,15 +111,11 @@ impl Cmd {
             let end = start + Duration::days(1);
             db.range(start, end).await?
         };
-        compute_stats(&history, self.count)?;
+        compute_stats(settings, &history, self.count)?;
         Ok(())
     }
 }
 
-// TODO: make this configurable?
-static COMMON_COMMAND_PREFIX: &[&str] = &["sudo"];
-static COMMON_SUBCOMMAND_PREFIX: &[&str] = &["cargo", "go", "git", "npm", "yarn", "pnpm"];
-
 fn first_non_whitespace(s: &str) -> Option<usize> {
     s.char_indices()
         // find the first non whitespace char
@@ -134,7 +132,7 @@ fn first_whitespace(s: &str) -> usize {
         .map_or(s.len(), |(i, _)| i)
 }
 
-fn interesting_command(mut command: &str) -> &str {
+fn interesting_command<'a>(settings: &Settings, mut command: &'a str) -> &'a str {
     // compute command prefix
     // we loop here because we might be working with a common command prefix (eg sudo) that we want to trim off
     let (i, prefix) = loop {
@@ -142,7 +140,7 @@ fn interesting_command(mut command: &str) -> &str {
         let prefix = &command[..i];
 
         // is it a common prefix
-        if COMMON_COMMAND_PREFIX.contains(&prefix) {
+        if settings.stats.common_prefix.contains(&String::from(prefix)) {
             command = command[i..].trim_start();
             if command.is_empty() {
                 // no commands following, just use the prefix
@@ -164,7 +162,14 @@ fn interesting_command(mut command: &str) -> &str {
 
     match subcommand_indices {
         // if there is a subcommand and it's a common one, then count the full prefix + subcommand
-        Some(end) if COMMON_SUBCOMMAND_PREFIX.contains(&prefix) => &command[..end],
+        Some(end)
+            if settings
+                .stats
+                .common_subcommands
+                .contains(&String::from(prefix)) =>
+        {
+            &command[..end]
+        }
         // otherwise just count the main command
         _ => prefix,
     }
@@ -172,16 +177,23 @@ fn interesting_command(mut command: &str) -> &str {
 
 #[cfg(test)]
 mod tests {
+    use atuin_client::settings::Settings;
+
     use super::interesting_command;
 
     #[test]
     fn interesting_commands() {
-        assert_eq!(interesting_command("cargo"), "cargo");
-        assert_eq!(interesting_command("cargo build foo bar"), "cargo build");
+        let settings = Settings::default();
+
+        assert_eq!(interesting_command(&settings, "cargo"), "cargo");
+        assert_eq!(
+            interesting_command(&settings, "cargo build foo bar"),
+            "cargo build"
+        );
         assert_eq!(
-            interesting_command("sudo   cargo build foo bar"),
+            interesting_command(&settings, "sudo   cargo build foo bar"),
             "cargo build"
         );
-        assert_eq!(interesting_command("sudo"), "sudo");
+        assert_eq!(interesting_command(&settings, "sudo"), "sudo");
     }
 }