From 3daf3ddf26ab8f2860e2c961bbe8b8d691049d08 Mon Sep 17 00:00:00 2001
From: Kevin Song <chipbuster@users.noreply.github.com>
Date: Thu, 8 Aug 2019 10:25:30 -0700
Subject: [PATCH] feat: implement timer module (#118)

Implement a timer module that takes a commandline argument, the number of seconds the last job took to complete, and displays it if appropriate.

Alters shell initialization files to compute this number using date +%s where needed.

Adds a config section to configure minimum amount of time before timer is shown (default is 2s)
---
 docs/config/README.md           | 28 ++++++++++
 src/init.rs                     | 81 ++++++++++++++++++++++++-----
 src/main.rs                     | 13 ++++-
 src/modules/cmd_duration.rs     | 90 +++++++++++++++++++++++++++++++++
 src/modules/mod.rs              |  2 +
 src/print.rs                    |  1 +
 tests/testsuite/cmd_duration.rs | 80 +++++++++++++++++++++++++++++
 tests/testsuite/main.rs         |  1 +
 8 files changed, 280 insertions(+), 16 deletions(-)
 create mode 100644 src/modules/cmd_duration.rs
 create mode 100644 tests/testsuite/cmd_duration.rs

diff --git a/docs/config/README.md b/docs/config/README.md
index 2ed8bd7e..772fa6b4 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -109,6 +109,33 @@ command had an unsuccessful status code (non-zero).
 symbol = "❯"
 ```
 
+## Command Duration
+
+The `cmd_duration` module shows how long the last command took to execute.
+The module will be shown only if the command took longer than two seconds, or
+the `min_time` config value, if it exists.
+
+::: warning NOTE
+Command duration is currently not supported in `bash`. See
+[this issue](https://github.com/starship/starship/issues/124) for more details.
+:::
+
+### Options
+
+| Variable   | Default | Description                         |
+| ---------- | ------- | ----------------------------------- |
+| `min_time` | `2`     | Shortest duration to show time for. |
+| `disabled` | `false` | Disables the `cmd_duration` module. |
+
+### Example
+
+```toml
+# ~/.config/starship.toml
+
+[cmd_duration]
+min_time = 4
+```
+
 ## Directory
 
 The `directory` module shows the path to your current directory, truncated to
@@ -366,3 +393,4 @@ The module will be shown if any of the following conditions are met:
 [username]
 disabled = true
 ```
+
diff --git a/src/init.rs b/src/init.rs
index 396ba72d..3941e92d 100644
--- a/src/init.rs
+++ b/src/init.rs
@@ -1,33 +1,26 @@
 use std::ffi::OsStr;
 use std::path::Path;
 
+/* We need to send execution time to the prompt for the cmd_duration module. For fish,
+this is fairly straightforward. For bash and zsh, we'll need to use several
+shell utilities to get the time, as well as render the prompt */
+
 pub fn init(shell_name: &str) {
     log::debug!("Shell name: {}", shell_name);
 
     let shell_basename = Path::new(shell_name).file_stem().and_then(OsStr::to_str);
 
     let setup_script = match shell_basename {
-        // The contents of `PROMPT_COMMAND` are executed as a regular Bash command
-        // just before Bash displays a prompt.
         Some("bash") => {
-            let script = "
-            PROMPT_COMMAND=starship_prompt
-            
-            starship_prompt() {
-                PS1=\"$(starship prompt --status=$?)\"
-            }";
+            let script = BASH_INIT;
             Some(script)
         }
-        // `precmd` executes a command before the zsh prompt is displayed.
         Some("zsh") => {
-            let script = "
-            precmd() {
-                PROMPT=\"$(starship prompt --status=$?)\"
-            }";
+            let script = ZSH_INIT;
             Some(script)
         }
         Some("fish") => {
-            let script = "function fish_prompt; starship prompt --status=$status; end";
+            let script = FISH_INIT;
             Some(script)
         }
         None => {
@@ -58,3 +51,63 @@ pub fn init(shell_name: &str) {
         print!("{}", script);
     }
 }
+
+/* Bash does not currently support command durations (see issue #124) for details
+https://github.com/starship/starship/issues/124
+*/
+
+const BASH_INIT: &str = r##"
+starship_precmd() {
+        PS1="$(starship prompt --status=$?)";
+};
+PROMPT_COMMAND=starship_precmd;
+"##;
+/* TODO: Once warning/error system is implemented in starship, print a warning
+if starship will not be printing timing due to DEBUG clobber error */
+
+/* For zsh: preexec_functions and precmd_functions provide preexec/precmd in a
+   way that lets us avoid clobbering them.
+
+   Zsh quirk: preexec() is only fired if a command is actually run (unlike in
+   bash, where spamming empty commands still triggers DEBUG). This means a user
+   spamming ENTER at an empty command line will see increasing runtime (since
+   preexec never actually fires to reset the start time).
+
+   To fix this, only pass the time if STARSHIP_START_TIME is defined, and unset
+   it after passing the time, so that we only measure actual commands.
+*/
+
+const ZSH_INIT: &str = r##"
+starship_precmd() {
+    STATUS=$?;
+    if [[ $STARSHIP_START_TIME ]]; then
+        STARSHIP_END_TIME="$(date +%s)";
+        STARSHIP_DURATION=$((STARSHIP_END_TIME - STARSHIP_START_TIME));
+        PROMPT="$(starship prompt --status=$STATUS --cmd-duration=$STARSHIP_DURATION)";
+        unset STARSHIP_START_TIME;
+    else
+        PROMPT="$(starship prompt --status=$STATUS)";
+    fi
+};
+starship_preexec(){
+    STARSHIP_START_TIME="$(date +%s)"
+};
+if [[ ${precmd_functions[(ie)starship_precmd]} -gt ${#precmd_functions} ]]; then
+    precmd_functions+=(starship_precmd);
+fi;
+if [[ ${preexec_functions[(ie)starship_preexec]} -gt ${#preexec_functions} ]]; then
+    preexec_functions+=(starship_preexec);
+fi;
+STARSHIP_START_TIME="$(date +%s)";
+"##;
+
+/* Fish setup is simple because they give us CMD_DURATION. Just account for name
+changes between 2.7/3.0 and do some math to convert ms->s and we can use it */
+const FISH_INIT: &str = r##"
+function fish_prompt;
+    set -l exit_code $status;
+    set -l CMD_DURATION "$CMD_DURATION$cmd_duration";
+    set -l starship_duration (math --scale=0 "$CMD_DURATION / 1000");
+    starship prompt --status=$exit_code --cmd-duration=$starship_duration;
+end;
+"##;
diff --git a/src/main.rs b/src/main.rs
index a7ac01a6..fa6f1488 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -36,6 +36,13 @@ fn main() {
         )
         .required(true);
 
+    let cmd_duration_arg = Arg::with_name("cmd_duration")
+        .short("d")
+        .long("cmd-duration")
+        .value_name("CMD_DURATION")
+        .help("The execution duration of the last command, in seconds")
+        .takes_value(true);
+
     let matches = App::new("starship")
         .about("The cross-shell prompt for astronauts. ☄🌌️")
         // pull the version number from Cargo.toml
@@ -53,7 +60,8 @@ fn main() {
             SubCommand::with_name("prompt")
                 .about("Prints the full starship prompt")
                 .arg(&status_code_arg)
-                .arg(&path_arg),
+                .arg(&path_arg)
+                .arg(&cmd_duration_arg),
         )
         .subcommand(
             SubCommand::with_name("module")
@@ -64,7 +72,8 @@ fn main() {
                         .required(true),
                 )
                 .arg(&status_code_arg)
-                .arg(&path_arg),
+                .arg(&path_arg)
+                .arg(&cmd_duration_arg),
         )
         .get_matches();
 
diff --git a/src/modules/cmd_duration.rs b/src/modules/cmd_duration.rs
new file mode 100644
index 00000000..1f0cd3e4
--- /dev/null
+++ b/src/modules/cmd_duration.rs
@@ -0,0 +1,90 @@
+use crate::config::Config;
+use ansi_term::Color;
+
+use super::{Context, Module};
+
+/// Outputs the time it took the last command to execute
+///
+/// Will only print if last command took more than a certain amount of time to
+/// execute. Default is two seconds, but can be set by config option `min_time`.
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("cmd_duration")?;
+
+    let arguments = &context.arguments;
+    let elapsed = arguments
+        .value_of("cmd_duration")
+        .unwrap_or("invalid_time")
+        .parse::<u64>()
+        .ok()?;
+
+    let signed_config_min = module.config_value_i64("min_time").unwrap_or(2);
+
+    /* TODO: Once error handling is implemented, warn the user if their config
+    min time is nonsensical */
+    if signed_config_min < 0 {
+        log::debug!(
+            "[WARN]: min_time in [cmd_duration] ({}) was less than zero",
+            signed_config_min
+        );
+        return None;
+    }
+
+    let config_min = signed_config_min as u64;
+
+    let module_color = match elapsed {
+        time if time < config_min => return None,
+        _ => Color::Yellow.bold(),
+    };
+
+    module.set_style(module_color);
+    module.new_segment("cmd_duration", &format!("took {}", render_time(elapsed)));
+    module.get_prefix().set_value("");
+
+    Some(module)
+}
+
+// Render the time into a nice human-readable string
+fn render_time(raw_seconds: u64) -> String {
+    // Calculate a simple breakdown into days/hours/minutes/seconds
+    let (seconds, raw_minutes) = (raw_seconds % 60, raw_seconds / 60);
+    let (minutes, raw_hours) = (raw_minutes % 60, raw_minutes / 60);
+    let (hours, days) = (raw_hours % 24, raw_hours / 24);
+
+    let components = [days, hours, minutes, seconds];
+    let suffixes = ["d", "h", "m", "s"];
+
+    let rendered_components: Vec<String> = components
+        .iter()
+        .zip(&suffixes)
+        .map(render_time_component)
+        .collect();
+    rendered_components.join("")
+}
+
+/// Render a single component of the time string, giving an empty string if component is zero
+fn render_time_component((component, suffix): (&u64, &&str)) -> String {
+    match component {
+        0 => String::new(),
+        n => format!("{}{}", n, suffix),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_10s() {
+        assert_eq!(render_time(10 as u64), "10s")
+    }
+    fn test_90s() {
+        assert_eq!(render_time(90 as u64), "1m30s")
+    }
+    fn test_10110s() {
+        assert_eq!(render_time(10110 as u64), "1h48m30s")
+    }
+    fn test_1d() {
+        assert_eq!(render_time(86400 as u64), "1d")
+    }
+
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index ffa21a81..5f3fcbba 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,5 +1,6 @@
 mod battery;
 mod character;
+mod cmd_duration;
 mod directory;
 mod git_branch;
 mod git_status;
@@ -28,6 +29,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
         "git_status" => git_status::module(context),
         "username" => username::module(context),
         "battery" => battery::module(context),
+        "cmd_duration" => cmd_duration::module(context),
 
         _ => panic!("Unknown module: {}", module),
     }
diff --git a/src/print.rs b/src/print.rs
index 3cd185e4..c3d46f5d 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -18,6 +18,7 @@ const PROMPT_ORDER: &[&str] = &[
     "rust",
     "python",
     "go",
+    "cmd_duration",
     "line_break",
     "character",
 ];
diff --git a/tests/testsuite/cmd_duration.rs b/tests/testsuite/cmd_duration.rs
new file mode 100644
index 00000000..a7cf5e95
--- /dev/null
+++ b/tests/testsuite/cmd_duration.rs
@@ -0,0 +1,80 @@
+use ansi_term::Color;
+use std::fs;
+use std::io;
+use std::path::Path;
+use tempfile::TempDir;
+
+use crate::common::{self, TestCommand};
+
+#[test]
+fn config_blank_duration_1s() -> io::Result<()> {
+    let output = common::render_module("cmd_duration")
+        .arg("--cmd-duration=1")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = "";
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn config_blank_duration_5s() -> io::Result<()> {
+    let output = common::render_module("cmd_duration")
+        .arg("--cmd-duration=5")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("{} ", Color::Yellow.bold().paint("took 5s"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn config_5s_duration_3s() -> io::Result<()> {
+    let output = common::render_module("cmd_duration")
+        .use_config(toml::toml! {
+            [cmd_duration]
+            min_time = 5
+        })
+        .arg("--cmd-duration=3")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = "";
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn config_5s_duration_10s() -> io::Result<()> {
+    let output = common::render_module("cmd_duration")
+        .use_config(toml::toml! {
+            [cmd_duration]
+            min_time = 5
+        })
+        .arg("--cmd-duration=10")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("{} ", Color::Yellow.bold().paint("took 10s"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn config_disabled() -> io::Result<()> {
+    let output = common::render_module("cmd_duration")
+        .use_config(toml::toml! {
+            [cmd_duration]
+            disabled = true
+            min_time = 5
+        })
+        .arg("--cmd-duration=10")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = "";
+    assert_eq!(expected, actual);
+    Ok(())
+}
diff --git a/tests/testsuite/main.rs b/tests/testsuite/main.rs
index d7989e7a..ab520117 100644
--- a/tests/testsuite/main.rs
+++ b/tests/testsuite/main.rs
@@ -1,4 +1,5 @@
 mod character;
+mod cmd_duration;
 mod common;
 mod configuration;
 mod directory;