feat(nix): support new `nix shell` command (#4724)

* Support `nix shell`

* Remove unnecessary `Debug` implementation

* Add test to detect false positive

* Improve detection of `/nix/store` in $PATH

* Add docs about unknown state

* Gate under `heuristic` flag

* Regenerate config schema

.github/config-schema.json

@@ -989,10 +989,12 @@
       "default": {
         "disabled": false,
         "format": "via [$symbol$state( \\($name\\))]($style) ",
+        "heuristic": false,
         "impure_msg": "impure",
         "pure_msg": "pure",
         "style": "bold blue",
-        "symbol": "❄️  "
+        "symbol": "❄️  ",
+        "unknown_msg": ""
       },
       "allOf": [
         {
@@ -4066,9 +4068,17 @@
           "default": "pure",
           "type": "string"
         },
+        "unknown_msg": {
+          "default": "",
+          "type": "string"
+        },
         "disabled": {
           "default": false,
           "type": "boolean"
+        },
+        "heuristic": {
+          "default": false,
+          "type": "boolean"
         }
       },
       "additionalProperties": false

docs/config/README.md

@@ -2719,13 +2719,15 @@ The module will be shown when inside a nix-shell environment.
 ### Options
 
 | Option        | Default                                      | Description                                                           |
-| ------------ | -------------------------------------------- | ----------------------------------------------------- |
+| ------------- | -------------------------------------------- | --------------------------------------------------------------------- |
 | `format`      | `'via [$symbol$state( \($name\))]($style) '` | The format for the module.                                            |
 | `symbol`      | `'❄️ '`                                       | A format string representing the symbol of nix-shell.                 |
 | `style`       | `'bold blue'`                                | The style for the module.                                             |
 | `impure_msg`  | `'impure'`                                   | A format string shown when the shell is impure.                       |
 | `pure_msg`    | `'pure'`                                     | A format string shown when the shell is pure.                         |
+| `unknown_msg` | `''`                                         | A format string shown when it is unknown if the shell is pure/impure. |
 | `disabled`    | `false`                                      | Disables the `nix_shell` module.                                      |
+| `heuristic`   | `false`                                      | Attempts to detect new `nix shell`-style shells with a heuristic.     |
 
 ### Variables
 
@@ -2747,6 +2749,7 @@ The module will be shown when inside a nix-shell environment.
 disabled = true
 impure_msg = '[impure shell](bold red)'
 pure_msg = '[pure shell](bold green)'
+unknown_msg = '[unknown shell](bold yellow)'
 format = 'via [☃️ $state( \($name\))](bold blue) '
 ```
 

src/configs/nix_shell.rs

@@ -13,7 +13,9 @@ pub struct NixShellConfig<'a> {
     pub style: &'a str,
     pub impure_msg: &'a str,
     pub pure_msg: &'a str,
+    pub unknown_msg: &'a str,
     pub disabled: bool,
+    pub heuristic: bool,
 }
 
 /* The trailing double spaces in `symbol` are needed to work around issues with
@@ -27,7 +29,9 @@ impl<'a> Default for NixShellConfig<'a> {
             style: "bold blue",
             impure_msg: "impure",
             pure_msg: "pure",
+            unknown_msg: "",
             disabled: false,
+            heuristic: false,
         }
     }
 }

src/modules/nix_shell.rs

@@ -3,6 +3,43 @@ use super::{Context, Module, ModuleConfig};
 use crate::configs::nix_shell::NixShellConfig;
 use crate::formatter::StringFormatter;
 
+enum NixShellType {
+    Pure,
+    Impure,
+    /// We're in a Nix shell, but we don't know which type.
+    /// This can only happen in a `nix shell` shell (not a `nix-shell` one).
+    Unknown,
+}
+
+impl NixShellType {
+    fn detect_shell_type(use_heuristic: bool, context: &Context) -> Option<NixShellType> {
+        use NixShellType::*;
+
+        let shell_type = context.get_env("IN_NIX_SHELL");
+        match shell_type.as_deref() {
+            Some("pure") => return Some(Pure),
+            Some("impure") => return Some(Impure),
+            _ => {}
+        };
+
+        if use_heuristic {
+            Self::in_new_nix_shell(context).map(|_| Unknown)
+        } else {
+            None
+        }
+    }
+
+    // Hack to detect if we're in a `nix shell` (in constrast to a `nix-shell`).
+    // A better way to do this will be enabled by https://github.com/NixOS/nix/issues/6677.
+    fn in_new_nix_shell(context: &Context) -> Option<()> {
+        let path = context.get_env("PATH")?;
+
+        std::env::split_paths(&path)
+            .any(|path| path.starts_with("/nix/store"))
+            .then_some(())
+    }
+}
+
 /// Creates a module showing if inside a nix-shell
 ///
 /// The module will use the `$IN_NIX_SHELL` and `$name` environment variable to
@@ -11,24 +48,25 @@ use crate::formatter::StringFormatter;
 /// The following options are availables:
 ///     - `impure_msg` (string)  // change the impure msg
 ///     - `pure_msg` (string)    // change the pure msg
+///     - `unknown_msg` (string) // change the unknown message
 ///
 /// Will display the following:
 ///     - pure (name)    // $name == "name" in a pure nix-shell
 ///     - impure (name)  // $name == "name" in an impure nix-shell
 ///     - pure           // $name == "" in a pure nix-shell
 ///     - impure         // $name == "" in an impure nix-shell
+///     - unknown (name) // $name == "name" in an unknown nix-shell
+///     - unknown        // $name == "" in an unknown nix-shell
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let mut module = context.new_module("nix_shell");
     let config: NixShellConfig = NixShellConfig::try_load(module.config);
 
     let shell_name = context.get_env("name");
-    let shell_type = context.get_env("IN_NIX_SHELL")?;
-    let shell_type_format = match shell_type.as_ref() {
-        "impure" => config.impure_msg,
-        "pure" => config.pure_msg,
-        _ => {
-            return None;
-        }
+    let shell_type = NixShellType::detect_shell_type(config.heuristic, context)?;
+    let shell_type_format = match shell_type {
+        NixShellType::Pure => config.pure_msg,
+        NixShellType::Impure => config.impure_msg,
+        NixShellType::Unknown => config.unknown_msg,
     };
 
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
@@ -130,4 +168,62 @@ mod tests {
 
         assert_eq!(expected, actual);
     }
+
+    #[test]
+    fn new_nix_shell() {
+        let actual = ModuleRenderer::new("nix_shell")
+            .env(
+                "PATH",
+                "/nix/store/v7qvqv81jp0cajvrxr9x072jgqc01yhi-nix-info/bin:/Users/user/.cargo/bin",
+            )
+            .config(toml::toml! {
+                [nix_shell]
+                heuristic = true
+            })
+            .collect();
+        let expected = Some(format!("via {} ", Color::Blue.bold().paint("❄️  ")));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn no_new_nix_shell() {
+        let actual = ModuleRenderer::new("nix_shell")
+            .env("PATH", "/Users/user/.cargo/bin")
+            .config(toml::toml! {
+                [nix_shell]
+                heuristic = true
+            })
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn no_new_nix_shell_with_nix_store_subdirectory() {
+        let actual = ModuleRenderer::new("nix_shell")
+            .env("PATH", "/Users/user/some/nix/store/subdirectory")
+            .config(toml::toml! {
+                [nix_shell]
+                heuristic = true
+            })
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn no_new_nix_shell_when_heuristic_is_disabled() {
+        let actual = ModuleRenderer::new("nix_shell")
+            .env(
+                "PATH",
+                "/nix/store/v7qvqv81jp0cajvrxr9x072jgqc01yhi-nix-info/bin:/Users/user/.cargo/bin",
+            )
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
 }