Add support for custom modules. (#916)

Cargo.lock

@@ -412,6 +412,14 @@ dependencies = [
  "unicode-normalization 0.1.12 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "indexmap"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "autocfg 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "itoa"
 version = "0.4.5"
@@ -1135,6 +1143,7 @@ name = "toml"
 version = "0.5.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 dependencies = [
+ "indexmap 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.105 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
@@ -1325,6 +1334,7 @@ dependencies = [
 "checksum http 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b708cc7f06493459026f53b9a61a7a121a5d1ec6238dee58ea4941132b30156b"
 "checksum humantime 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "df004cfca50ef23c36850aaaa59ad52cc70d0e90243c3c7737a4dd32dc7a3c4f"
 "checksum idna 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "02e2673c30ee86b5b96a9cb52ad15718aa1f966f5ab9ad54a8b95d5ca33120a9"
+"checksum indexmap 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "076f042c5b7b98f31d205f1249267e12a6518c1481e9dae9764af19b707d2292"
 "checksum itoa 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "b8b7a7c0c47db5545ed3fef7468ee7bb5b74691498139e4b3f6a20685dc6dd8e"
 "checksum jobserver 0.1.21 (registry+https://github.com/rust-lang/crates.io-index)" = "5c71313ebb9439f74b00d9d2dcec36440beaf57a6aa0623068441dd7cd81a7f2"
 "checksum kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "7507624b29483431c0ba2d82aece8ca6cdba9382bff4ddd0f7490560c056098d"

Cargo.toml

@@ -32,7 +32,7 @@ clap = "2.33.0"
 ansi_term = "0.12.1"
 dirs = "2.0.2"
 git2 = { version = "0.13.1", default-features = false, features = [] }
-toml = "0.5.6"
+toml = { version = "0.5.6", features = ["preserve_order"] }
 serde_json = "1.0.51"
 rayon = "1.3.0"
 pretty_env_logger = "0.4.0"

docs/config/README.md

@@ -122,6 +122,7 @@ prompt_order = [
     "env_var",
     "crystal",
     "cmd_duration",
+    "custom",
     "line_break",
     "jobs",
     "battery",
@@ -1299,3 +1300,56 @@ The module will be shown if any of the following conditions are met:
 [username]
 disabled = true
 ```
+
+## Custom commands
+
+The `custom` modules show the output of some arbitrary commands.
+
+These modules will be shown if any of the following conditions are met:
+- The current directory contains a file whose name is in `files`
+- The current directory contains a directory whose name is in `directories`
+- The current directory contains a file whose extension is in `extensions`
+- The `when` command returns 0
+
+::: tip
+
+Multiple custom modules can be defined by using a `.`.
+
+:::
+
+::: tip
+
+The order in which custom modules are shown can be individually set
+by setting `custom.foo` in `prompt_order`. By default, the `custom` module
+will simply show all custom modules in the order they were defined.
+
+:::
+
+### Options
+
+| Variable      | Default             | Description                                                                  |
+| ------------- | ------------------- | ---------------------------------------------------------------------------- |
+| `command`     |                     | The command whose output should be printed.                                  |
+| `when`        |                     | A shell command used as a condition to show the module. The module will be shown if the command returns a `0` status code. |
+| `shell`       |                     | The path to the shell to use to execute the command. If unset, it will fallback to STARSHIP_SHELL and then to "sh". |
+| `description` | `"<custom module>"` | The description of the module that is shown when running `starship explain`. |
+| `files`       | `[]`                | The files that will be searched in the working directory for a match.        |
+| `directories` | `[]`                | The directories that will be searched in the working directory for a match.  |
+| `extensions`  | `[]`                | The extensions that will be searched in the working directory for a match.   |
+| `symbol`      | `""`                | The symbol used before displaying the command output.                        |
+| `style`       | `"bold green"`      | The style for the module.                                                    |
+| `prefix`      | `""`                | Prefix to display immediately before the command output.                     |
+| `suffix`      | `""`                | Suffix to display immediately after the command output.                      |
+| `disabled`    | `false`             | Disables this `custom` module.                                               |
+
+### Example
+
+```toml
+# ~/.config/starship.toml
+
+[custom.foo]
+command = "echo foo"  # shows output of command
+files = ["foo"]       # can specify filters
+when = """ test "$HOME" == "$PWD" """
+prefix = " transcending "
+```

src/config.rs

@@ -218,6 +218,26 @@ impl StarshipConfig {
         module_config
     }
 
+    /// Get the subset of the table for a custom module by its name
+    pub fn get_custom_module_config(&self, module_name: &str) -> Option<&Value> {
+        let module_config = self.get_custom_modules()?.get(module_name);
+        if module_config.is_some() {
+            log::debug!(
+                "Custom config found for \"{}\": \n{:?}",
+                &module_name,
+                &module_config
+            );
+        } else {
+            log::trace!("No custom config found for \"{}\"", &module_name);
+        }
+        module_config
+    }
+
+    /// Get the table of all the registered custom modules, if any
+    pub fn get_custom_modules(&self) -> Option<&toml::value::Table> {
+        self.config.as_ref()?.as_table()?.get("custom")?.as_table()
+    }
+
     pub fn get_root_config(&self) -> StarshipRootConfig {
         if let Some(root_config) = &self.config {
             StarshipRootConfig::load(root_config)

src/configs/custom.rs

@@ -0,0 +1,96 @@
+use crate::config::{ModuleConfig, RootModuleConfig, SegmentConfig};
+
+use ansi_term::Style;
+use starship_module_config_derive::ModuleConfig;
+
+#[derive(Clone, Default, PartialEq)]
+pub struct Files<'a>(pub Vec<&'a str>);
+
+#[derive(Clone, Default, PartialEq)]
+pub struct Extensions<'a>(pub Vec<&'a str>);
+
+#[derive(Clone, Default, PartialEq)]
+pub struct Directories<'a>(pub Vec<&'a str>);
+
+#[derive(Clone, ModuleConfig)]
+pub struct CustomConfig<'a> {
+    pub symbol: Option<SegmentConfig<'a>>,
+    pub command: &'a str,
+    pub when: Option<&'a str>,
+    pub shell: Option<&'a str>,
+    pub description: &'a str,
+    pub style: Option<Style>,
+    pub disabled: bool,
+    pub prefix: Option<&'a str>,
+    pub suffix: Option<&'a str>,
+    pub files: Files<'a>,
+    pub extensions: Extensions<'a>,
+    pub directories: Directories<'a>,
+}
+
+impl<'a> RootModuleConfig<'a> for CustomConfig<'a> {
+    fn new() -> Self {
+        CustomConfig {
+            symbol: None,
+            command: "",
+            when: None,
+            shell: None,
+            description: "<custom config>",
+            style: None,
+            disabled: false,
+            prefix: None,
+            suffix: None,
+            files: Files::default(),
+            extensions: Extensions::default(),
+            directories: Directories::default(),
+        }
+    }
+}
+
+impl<'a> ModuleConfig<'a> for Files<'a> {
+    fn from_config(config: &'a toml::Value) -> Option<Self> {
+        let mut files = Vec::new();
+
+        for item in config.as_array()? {
+            if let Some(file) = item.as_str() {
+                files.push(file);
+            } else {
+                log::debug!("Unexpected file {:?}", item);
+            }
+        }
+
+        Some(Files(files))
+    }
+}
+
+impl<'a> ModuleConfig<'a> for Extensions<'a> {
+    fn from_config(config: &'a toml::Value) -> Option<Self> {
+        let mut extensions = Vec::new();
+
+        for item in config.as_array()? {
+            if let Some(file) = item.as_str() {
+                extensions.push(file);
+            } else {
+                log::debug!("Unexpected extension {:?}", item);
+            }
+        }
+
+        Some(Extensions(extensions))
+    }
+}
+
+impl<'a> ModuleConfig<'a> for Directories<'a> {
+    fn from_config(config: &'a toml::Value) -> Option<Self> {
+        let mut directories = Vec::new();
+
+        for item in config.as_array()? {
+            if let Some(file) = item.as_str() {
+                directories.push(file);
+            } else {
+                log::debug!("Unexpected directory {:?}", item);
+            }
+        }
+
+        Some(Directories(directories))
+    }
+}

src/configs/mod.rs

@@ -4,6 +4,7 @@ pub mod character;
 pub mod cmd_duration;
 pub mod conda;
 pub mod crystal;
+pub mod custom;
 pub mod directory;
 pub mod docker_context;
 pub mod dotnet;

src/configs/starship_root.rs

@@ -51,6 +51,7 @@ impl<'a> RootModuleConfig<'a> for StarshipRootConfig<'a> {
                 "aws",
                 "env_var",
                 "cmd_duration",
+                "custom",
                 "line_break",
                 "jobs",
                 #[cfg(feature = "battery")]

src/context.rs

@@ -112,6 +112,15 @@ impl<'a> Context<'a> {
         disabled == Some(true)
     }
 
+    /// Return whether the specified custom module has a `disabled` option set to true.
+    /// If it doesn't exist, `None` is returned.
+    pub fn is_custom_module_disabled_in_config(&self, name: &str) -> Option<bool> {
+        let config = self.config.get_custom_module_config(name)?;
+        let disabled = Some(config).and_then(|table| table.as_table()?.get("disabled")?.as_bool());
+
+        Some(disabled == Some(true))
+    }
+
     // returns a new ScanDir struct with reference to current dir_files of context
     // see ScanDir for methods
     pub fn try_begin_scan(&'a self) -> Option<ScanDir<'a>> {

src/modules/custom.rs

@@ -0,0 +1,288 @@
+use ansi_term::Color;
+use std::io::Write;
+use std::process::{Command, Output, Stdio};
+
+use super::{Context, Module, RootModuleConfig};
+
+use crate::{config::SegmentConfig, configs::custom::CustomConfig};
+
+/// Creates a custom module with some configuration
+///
+/// The relevant TOML config will set the files, extensions, and directories needed
+/// for the module to be displayed. If none of them match, and optional "when"
+/// command can be run -- if its result is 0, the module will be shown.
+///
+/// Finally, the content of the module itself is also set by a command.
+pub fn module<'a>(name: &'a str, context: &'a Context) -> Option<Module<'a>> {
+    let toml_config = context.config.get_custom_module_config(name).expect(
+        "modules::custom::module should only be called after ensuring that the module exists",
+    );
+    let config = CustomConfig::load(toml_config);
+
+    let mut scan_dir = context.try_begin_scan()?;
+
+    if !config.files.0.is_empty() {
+        scan_dir = scan_dir.set_files(&config.files.0);
+    }
+    if !config.extensions.0.is_empty() {
+        scan_dir = scan_dir.set_extensions(&config.extensions.0);
+    }
+    if !config.directories.0.is_empty() {
+        scan_dir = scan_dir.set_folders(&config.directories.0);
+    }
+
+    let mut is_match = scan_dir.is_match();
+
+    if !is_match {
+        if let Some(when) = config.when {
+            is_match = exec_when(when, config.shell);
+        }
+
+        if !is_match {
+            return None;
+        }
+    }
+
+    let mut module = Module::new(name, config.description, Some(toml_config));
+    let style = config.style.unwrap_or_else(|| Color::Green.bold());
+
+    if let Some(prefix) = config.prefix {
+        module.get_prefix().set_value(prefix);
+    }
+    if let Some(suffix) = config.suffix {
+        module.get_suffix().set_value(suffix);
+    }
+
+    if let Some(symbol) = config.symbol {
+        module.create_segment("symbol", &symbol);
+    }
+
+    if let Some(output) = exec_command(config.command, config.shell) {
+        let trimmed = output.trim();
+
+        if trimmed.is_empty() {
+            return None;
+        }
+
+        module.create_segment(
+            "output",
+            &SegmentConfig::new(&trimmed).with_style(Some(style)),
+        );
+
+        Some(module)
+    } else {
+        None
+    }
+}
+
+/// Return the invoking shell, using `shell` and fallbacking in order to STARSHIP_SHELL and "sh"
+#[cfg(not(windows))]
+fn get_shell(shell: Option<&str>) -> std::borrow::Cow<str> {
+    if let Some(forced_shell) = shell {
+        forced_shell.into()
+    } else if let Ok(env_shell) = std::env::var("STARSHIP_SHELL") {
+        env_shell.into()
+    } else {
+        "sh".into()
+    }
+}
+
+/// Attempt to run the given command in a shell by passing it as `stdin` to `get_shell()`
+#[cfg(not(windows))]
+fn shell_command(cmd: &str, shell: Option<&str>) -> Option<Output> {
+    let command = Command::new(get_shell(shell).as_ref())
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn();
+
+    let mut child = match command {
+        Ok(command) => command,
+        Err(_) => {
+            log::debug!(
+                "Could not launch command with given shell or STARSHIP_SHELL env variable, retrying with /bin/env sh"
+            );
+
+            Command::new("/bin/env")
+                .arg("sh")
+                .stdin(Stdio::piped())
+                .stdout(Stdio::piped())
+                .stderr(Stdio::piped())
+                .spawn()
+                .ok()?
+        }
+    };
+
+    child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
+    child.wait_with_output().ok()
+}
+
+/// Attempt to run the given command in a shell by passing it as `stdin` to `get_shell()`,
+/// or by invoking cmd.exe /C.
+#[cfg(windows)]
+fn shell_command(cmd: &str, shell: Option<&str>) -> Option<Output> {
+    let shell = if let Some(shell) = shell {
+        Some(std::borrow::Cow::Borrowed(shell))
+    } else if let Ok(env_shell) = std::env::var("STARSHIP_SHELL") {
+        Some(std::borrow::Cow::Owned(env_shell))
+    } else {
+        None
+    };
+
+    if let Some(forced_shell) = shell {
+        let command = Command::new(forced_shell.as_ref())
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .spawn();
+
+        if let Ok(mut child) = command {
+            child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
+
+            return child.wait_with_output().ok();
+        }
+
+        log::debug!(
+            "Could not launch command with given shell or STARSHIP_SHELL env variable, retrying with cmd.exe /C"
+        );
+    }
+
+    let command = Command::new("cmd.exe")
+        .arg("/C")
+        .arg(cmd)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn();
+
+    command.ok()?.wait_with_output().ok()
+}
+
+/// Execute the given command capturing all output, and return whether it return 0
+fn exec_when(cmd: &str, shell: Option<&str>) -> bool {
+    log::trace!("Running '{}'", cmd);
+
+    if let Some(output) = shell_command(cmd, shell) {
+        if !output.status.success() {
+            log::trace!("non-zero exit code '{:?}'", output.status.code());
+            log::trace!(
+                "stdout: {}",
+                std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
+            );
+            log::trace!(
+                "stderr: {}",
+                std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
+            );
+        }
+
+        output.status.success()
+    } else {
+        log::debug!("Cannot start command");
+
+        false
+    }
+}
+
+/// Execute the given command, returning its output on success
+fn exec_command(cmd: &str, shell: Option<&str>) -> Option<String> {
+    log::trace!("Running '{}'", cmd);
+
+    if let Some(output) = shell_command(cmd, shell) {
+        if !output.status.success() {
+            log::trace!("Non-zero exit code '{:?}'", output.status.code());
+            log::trace!(
+                "stdout: {}",
+                std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
+            );
+            log::trace!(
+                "stderr: {}",
+                std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
+            );
+            return None;
+        }
+
+        Some(String::from_utf8_lossy(&output.stdout).into())
+    } else {
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[cfg(not(windows))]
+    const SHELL: Option<&'static str> = Some("/bin/sh");
+    #[cfg(windows)]
+    const SHELL: Option<&'static str> = None;
+
+    #[cfg(not(windows))]
+    const FAILING_COMMAND: &str = "false";
+    #[cfg(windows)]
+    const FAILING_COMMAND: &str = "color 00";
+
+    const UNKNOWN_COMMAND: &str = "ydelsyiedsieudleylse dyesdesl";
+
+    #[test]
+    fn when_returns_right_value() {
+        assert!(exec_when("echo hello", SHELL));
+        assert!(!exec_when(FAILING_COMMAND, SHELL));
+    }
+
+    #[test]
+    fn when_returns_false_if_invalid_command() {
+        assert!(!exec_when(UNKNOWN_COMMAND, SHELL));
+    }
+
+    #[test]
+    #[cfg(not(windows))]
+    fn command_returns_right_string() {
+        assert_eq!(exec_command("echo hello", SHELL), Some("hello\n".into()));
+        assert_eq!(
+            exec_command("echo 강남스타일", SHELL),
+            Some("강남스타일\n".into())
+        );
+    }
+
+    #[test]
+    #[cfg(windows)]
+    fn command_returns_right_string() {
+        assert_eq!(exec_command("echo hello", SHELL), Some("hello\r\n".into()));
+        assert_eq!(
+            exec_command("echo 강남스타일", SHELL),
+            Some("강남스타일\r\n".into())
+        );
+    }
+
+    #[test]
+    #[cfg(not(windows))]
+    fn command_ignores_stderr() {
+        assert_eq!(
+            exec_command("echo foo 1>&2; echo bar", SHELL),
+            Some("bar\n".into())
+        );
+        assert_eq!(
+            exec_command("echo foo; echo bar 1>&2", SHELL),
+            Some("foo\n".into())
+        );
+    }
+
+    #[test]
+    #[cfg(windows)]
+    fn command_ignores_stderr() {
+        assert_eq!(
+            exec_command("echo foo 1>&2 & echo bar", SHELL),
+            Some("bar\r\n".into())
+        );
+        assert_eq!(
+            exec_command("echo foo& echo bar 1>&2", SHELL),
+            Some("foo\r\n".into())
+        );
+    }
+
+    #[test]
+    fn command_can_fail() {
+        assert_eq!(exec_command(FAILING_COMMAND, SHELL), None);
+        assert_eq!(exec_command(UNKNOWN_COMMAND, SHELL), None);
+    }
+}

src/modules/mod.rs

@@ -4,6 +4,7 @@ mod character;
 mod cmd_duration;
 mod conda;
 mod crystal;
+pub(crate) mod custom;
 mod directory;
 mod docker_context;
 mod dotnet;

src/print.rs

@@ -1,7 +1,7 @@
 use ansi_term::ANSIStrings;
 use clap::ArgMatches;
 use rayon::prelude::*;
-use std::fmt::Write as FmtWrite;
+use std::fmt::{self, Debug, Write as FmtWrite};
 use std::io::{self, Write};
 use unicode_width::UnicodeWidthChar;
 
@@ -132,12 +132,60 @@ pub fn explain(args: ArgMatches) {
 }
 
 fn compute_modules<'a>(context: &'a Context) -> Vec<Module<'a>> {
-    let mut prompt_order: Vec<&str> = Vec::new();
+    enum Mod<'a> {
+        Builtin(&'a str),
+        Custom(&'a str),
+    }
+
+    struct DebugCustomModules<'tmp>(&'tmp toml::value::Table);
+
+    impl Debug for DebugCustomModules<'_> {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            f.debug_list().entries(self.0.keys()).finish()
+        }
+    }
+
+    let mut prompt_order: Vec<Mod> = Vec::new();
 
     // Write out a custom prompt order
-    for module in context.config.get_root_config().prompt_order {
+    let config_prompt_order = context.config.get_root_config().prompt_order;
-        if ALL_MODULES.contains(&module) {
+
-            prompt_order.push(module);
+    for module in &config_prompt_order {
+        if ALL_MODULES.contains(module) {
+            // Write out a module if it isn't disabled
+            if !context.is_module_disabled_in_config(*module) {
+                prompt_order.push(Mod::Builtin(module));
+            }
+        } else if *module == "custom" {
+            // Write out all custom modules, except for those that are explicitly set
+            if let Some(custom_modules) = context.config.get_custom_modules() {
+                for (custom_module, config) in custom_modules {
+                    if should_add_implicit_custom_module(
+                        custom_module,
+                        config,
+                        &config_prompt_order,
+                    ) {
+                        prompt_order.push(Mod::Custom(custom_module));
+                    }
+                }
+            }
+        } else if module.starts_with("custom.") {
+            // Write out a custom module if it isn't disabled (and it exists...)
+            match context.is_custom_module_disabled_in_config(&module[7..]) {
+                Some(true) => (), // Module is disabled, we don't add it to the prompt
+                Some(false) => prompt_order.push(Mod::Custom(&module[7..])),
+                None => match context.config.get_custom_modules() {
+                    Some(modules) => log::debug!(
+                        "prompt_order contains custom module \"{}\", but no configuration was provided. Configuration for the following modules were provided: {:?}",
+                        module,
+                        DebugCustomModules(modules),
+                    ),
+                    None => log::debug!(
+                        "prompt_order contains custom module \"{}\", but no configuration was provided.",
+                        module,
+                    ),
+                },
+            }
         } else {
             log::debug!(
                 "Expected prompt_order to contain value from {:?}. Instead received {}",
@@ -149,12 +197,37 @@ fn compute_modules<'a>(context: &'a Context) -> Vec<Module<'a>> {
 
     prompt_order
         .par_iter()
-        .filter(|module| !context.is_module_disabled_in_config(module))
+        .map(|module| match module {
-        .map(|module| modules::handle(module, &context)) // Compute modules
+            Mod::Builtin(builtin) => modules::handle(builtin, context),
+            Mod::Custom(custom) => modules::custom::module(custom, context),
+        }) // Compute segments
         .flatten() // Remove segments set to `None`
         .collect::<Vec<Module<'a>>>()
 }
 
+fn should_add_implicit_custom_module(
+    custom_module: &str,
+    config: &toml::Value,
+    config_prompt_order: &[&str],
+) -> bool {
+    let is_explicitly_specified = config_prompt_order.iter().any(|x| {
+        x.len() == 7 + custom_module.len() && &x[..7] == "custom." && &x[7..] == custom_module
+    });
+
+    if is_explicitly_specified {
+        // The module is already specified explicitly, so we skip it
+        return false;
+    }
+
+    let false_value = toml::Value::Boolean(false);
+
+    !config
+        .get("disabled")
+        .unwrap_or(&false_value)
+        .as_bool()
+        .unwrap_or(false)
+}
+
 fn count_wide_chars(value: &str) -> usize {
     value.chars().filter(|c| c.width().unwrap_or(0) > 1).count()
 }