feat(directory): Add directory substitutions (#1183)

Adds an option to provide a table of strings to substitute in the
directory string. Fixes #1065.

Co-authored-by: Radu Butoi <butoi@google.com>

4 files changed, 91 insertions, 3 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index f481385d2..1d331f00d 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -399,9 +399,19 @@ it would have been `nixpkgs/pkgs`.
 
 | Variable                    | Default | Description                                                                              |
 | --------------------------- | ------- | ---------------------------------------------------------------------------------------- |
+| `substitutions`             |         | A table of substitutions to be made to the path.                                         |
 | `fish_style_pwd_dir_length` | `0`     | The number of characters to use when applying fish shell pwd path logic.                 |
 | `use_logical_path`          | `true`  | Displays the logical path provided by the shell (`PWD`) instead of the path from the OS. |
 
+`substitutions` allows you to define arbitrary replacements for literal strings that occur in the path, for example long network
+prefixes or development directories (i.e. Java). Note that this will disable the fish style PWD.
+
+```toml
+[directory.substitutions]
+"/Volumes/network/path" = "/net"
+"src/com/long/java/path" = "mypath"
+```
+
 `fish_style_pwd_dir_length` interacts with the standard truncation options in a way that can be surprising at first: if it's non-zero,
 the components of the path that would normally be truncated are instead displayed with that many characters. For example, the path
 `/built/this/city/on/rock/and/roll`, which would normally be displayed as as `rock/and/roll`, would be displayed as
diff --git a/src/configs/directory.rs b/src/configs/directory.rs
index 3d44b14dd..9c6aaf26e 100644
--- a/src/configs/directory.rs
+++ b/src/configs/directory.rs
@@ -1,4 +1,5 @@
 use crate::config::{ModuleConfig, RootModuleConfig};
+use std::collections::HashMap;
 
 use ansi_term::{Color, Style};
 use starship_module_config_derive::ModuleConfig;
@@ -7,6 +8,7 @@ use starship_module_config_derive::ModuleConfig;
 pub struct DirectoryConfig<'a> {
     pub truncation_length: i64,
     pub truncate_to_repo: bool,
+    pub substitutions: HashMap<String, &'a str>,
     pub fish_style_pwd_dir_length: i64,
     pub use_logical_path: bool,
     pub prefix: &'a str,
@@ -20,6 +22,7 @@ impl<'a> RootModuleConfig<'a> for DirectoryConfig<'a> {
             truncation_length: 3,
             truncate_to_repo: true,
             fish_style_pwd_dir_length: 0,
+            substitutions: HashMap::new(),
             use_logical_path: true,
             prefix: "in ",
             style: Color::Cyan.bold(),
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index 16cee2f39..464ef23b2 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -1,4 +1,5 @@
 use path_slash::PathExt;
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use unicode_segmentation::UnicodeSegmentation;
 
@@ -10,12 +11,15 @@ use crate::configs::directory::DirectoryConfig;
 
 /// Creates a module with the current directory
 ///
-/// Will perform path contraction and truncation.
+/// Will perform path contraction, substitution, and truncation.
 /// **Contraction**
 ///     - Paths beginning with the home directory or with a git repo right
 /// inside the home directory will be contracted to `~`
 ///     - Paths containing a git repo will contract to begin at the repo root
 ///
+/// **Substitution**
+/// Paths will undergo user-provided substitutions of substrings
+///
 /// **Truncation**
 /// Paths will be limited in length to `3` path components by default.
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
@@ -67,10 +71,14 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         _ => contract_path(current_dir, &home_dir, HOME_SYMBOL),
     };
 
+    let substituted_dir = substitute_path(dir_string, &config.substitutions);
+
     // Truncate the dir string to the maximum number of path components
-    let truncated_dir_string = truncate(dir_string, config.truncation_length as usize);
+    let truncated_dir_string = truncate(substituted_dir, config.truncation_length as usize);
 
-    if config.fish_style_pwd_dir_length > 0 {
+    // Substitutions could have changed the prefix, so don't allow them and
+    // fish-style path contraction together
+    if config.fish_style_pwd_dir_length > 0 && config.substitutions.is_empty() {
         // If user is using fish style path, we need to add the segment first
         let contracted_home_dir = contract_path(&current_dir, &home_dir, HOME_SYMBOL);
         let fish_style_dir = to_fish_style(
@@ -126,6 +134,18 @@ fn contract_path(full_path: &Path, top_level_path: &Path, top_level_replacement:
     )
 }
 
+/// Perform a list of string substitutions on the path
+///
+/// Given a list of (from, to) pairs, this will perform the string
+/// substitutions, in order, on the path. Any non-pair of strings is ignored.
+fn substitute_path(dir_string: String, substitutions: &HashMap<String, &str>) -> String {
+    let mut substituted_dir = dir_string;
+    for substitution_pair in substitutions.iter() {
+        substituted_dir = substituted_dir.replace(substitution_pair.0, substitution_pair.1);
+    }
+    substituted_dir
+}
+
 /// Takes part before contracted path and replaces it with fish style path
 ///
 /// Will take the first letter of each directory before the contracted path and
@@ -224,6 +244,17 @@ mod tests {
     }
 
     #[test]
+    fn substitute_prefix_and_middle() {
+        let full_path = "/absolute/path/foo/bar/baz";
+        let mut substitutions = HashMap::new();
+        substitutions.insert("/absolute/path".to_string(), "");
+        substitutions.insert("/bar/".to_string(), "/");
+
+        let output = substitute_path(full_path.to_string(), &substitutions);
+        assert_eq!(output, "/foo/baz");
+    }
+
+    #[test]
     fn fish_style_with_user_home_contracted_path() {
         let path = "~/starship/engines/booster/rocket";
         let output = to_fish_style(1, path.to_string(), "engines/booster/rocket");
diff --git a/tests/testsuite/directory.rs b/tests/testsuite/directory.rs
index c236dff80..afbf6121a 100644
--- a/tests/testsuite/directory.rs
+++ b/tests/testsuite/directory.rs
@@ -25,6 +25,50 @@ fn home_directory() -> io::Result<()> {
 }
 
 #[test]
+fn substituted_truncated_path() -> io::Result<()> {
+    let output = common::render_module("directory")
+        .arg("--path=/some/long/network/path/workspace/a/b/c/dev")
+        .use_config(toml::toml! {
+            [directory]
+            truncation_length = 4
+            [directory.substitutions]
+            "/some/long/network/path" = "/some/net"
+            "a/b/c" = "d"
+        })
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("in {} ", Color::Cyan.bold().paint("net/workspace/d/dev"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn strange_substitution() -> io::Result<()> {
+    let strange_sub = "/\\/;,!";
+    let output = common::render_module("directory")
+        .arg("--path=/foo/bar/regular/path")
+        .use_config(toml::toml! {
+            [directory]
+            truncation_length = 0
+            fish_style_pwd_dir_length = 2 // Overridden by substitutions
+            [directory.substitutions]
+            "regular" = strange_sub
+        })
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!(
+        "in {} ",
+        Color::Cyan
+            .bold()
+            .paint(format!("/foo/bar/{}/path", strange_sub))
+    );
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
 #[ignore]
 fn directory_in_home() -> io::Result<()> {
     let dir = home_dir().unwrap().join("starship/engine");