Add nix-shell support (#173)

7 files changed, 143 insertions, 0 deletions

diff --git a/README.md b/README.md
index 7357742aa..46e2bde4d 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,7 @@ The prompt shows information you need while you're working, while staying sleek
 - Current Ruby version (`💎`)
 - Current Python version (`🐍`)
 - Current Go version (`🐹`)
+- Nix-shell environment detection
 - Current version of package in current directory (`📦`)
   - npm (Node.js)
   - cargo (Rust)
diff --git a/docs/config/README.md b/docs/config/README.md
index 22f960987..acc91e77e 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -490,3 +490,29 @@ The module will be shown if any of the following conditions are met:
 [username]
 disabled = true
 ```
+
+## Nix-shell
+
+The `nix_shell` module shows the nix-shell environment.
+The module will be shown when inside a nix-shell environment.
+
+### Options
+
+| Variable     | Default  | Description                        |
+| ------------ | -------- | ---------------------------------- |
+| `disabled`   | `false`  | Disables the `username` module.    |
+| `use_name`   | `false`  | Display the name of the nix-shell. |
+| `impure_msg` | `impure` | Customize the "impure" msg.        |
+| `pure_msg`   | `pure`   | Customize the "pure" msg.          |
+
+### Example
+
+```toml
+# ~/.config/starship.toml
+
+[nix_shell]
+disabled = true
+use_name = true
+impure_msg = "impure shell"
+pure_msg = "pure shell"
+```
\ No newline at end of file
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 4aa67108e..46929c314 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -8,6 +8,7 @@ mod git_status;
 mod golang;
 mod jobs;
 mod line_break;
+mod nix_shell;
 mod nodejs;
 mod package;
 mod python;
@@ -35,6 +36,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
         "battery" => battery::module(context),
         "cmd_duration" => cmd_duration::module(context),
         "jobs" => jobs::module(context),
+        "nix_shell" => nix_shell::module(context),
 
         _ => {
             eprintln!("Error: Unknown module {}. Use starship module --list to list out all supported modules.", module);
diff --git a/src/modules/nix_shell.rs b/src/modules/nix_shell.rs
new file mode 100644
index 000000000..b37fcdcf4
--- /dev/null
+++ b/src/modules/nix_shell.rs
@@ -0,0 +1,54 @@
+use ansi_term::Color;
+use std::env;
+
+use super::{Context, Module};
+
+// IN_NIX_SHELL should be "pure" or "impure" but lorri uses "1" for "impure"
+// https://github.com/target/lorri/issues/140
+
+/// Creates a module showing if inside a nix-shell
+///
+/// The module will use the `$IN_NIX_SHELL` and `$name` environment variable to
+/// determine if it's inside a nix-shell and the name of it.
+///
+/// The following options are availables:
+///     - use_name   (bool)   // print the name of the nix-shell
+///     - impure_msg (string) // change the impure msg
+///     - pure_msg (string)   // change the pure msg
+///
+/// Will display the following:
+///     - name (pure)    // use_name == true in a pure nix-shell
+///     - name (impure)  // use_name == true in an impure nix-shell
+///     - pure           // use_name == false in a pure nix-shell
+///     - impure         // use_name == false in an impure nix-shell
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("nix_shell")?;
+
+    env::var("IN_NIX_SHELL")
+        .ok()
+        .and_then(|shell_type| {
+            if shell_type == "1" || shell_type == "impure" {
+                Some(module.config_value_str("impure_msg").unwrap_or("impure"))
+            } else if shell_type == "pure" {
+                Some(module.config_value_str("pure_msg").unwrap_or("pure"))
+            } else {
+                None
+            }
+        })
+        .map(|shell_type| {
+            if module.config_value_bool("use_name").unwrap_or(false) {
+                match env::var("name").ok() {
+                    Some(name) => format!("{} ({})", name, shell_type),
+                    None => shell_type.to_string(),
+                }
+            } else {
+                shell_type.to_string()
+            }
+        })
+        .map(|segment| {
+            let module_color = Color::Red.bold();
+            module.set_style(module_color);
+            module.new_segment("nix_shell", &segment);
+            module
+        })
+}
diff --git a/src/print.rs b/src/print.rs
index dc62dd94d..56ba2287f 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -22,6 +22,7 @@ const DEFAULT_PROMPT_ORDER: &[&str] = &[
     "rust",
     "python",
     "golang",
+    "nix_shell",
     "cmd_duration",
     "line_break",
     "jobs",
diff --git a/tests/testsuite/main.rs b/tests/testsuite/main.rs
index bceb14035..4ecde3792 100644
--- a/tests/testsuite/main.rs
+++ b/tests/testsuite/main.rs
@@ -7,6 +7,7 @@ mod golang;
 mod jobs;
 mod line_break;
 mod modules;
+mod nix_shell;
 mod nodejs;
 mod python;
 mod ruby;
diff --git a/tests/testsuite/nix_shell.rs b/tests/testsuite/nix_shell.rs
new file mode 100644
index 000000000..b278e3cde
--- /dev/null
+++ b/tests/testsuite/nix_shell.rs
@@ -0,0 +1,58 @@
+use ansi_term::Color;
+use std::io;
+
+use crate::common;
+
+#[test]
+fn no_env_variables() -> io::Result<()> {
+    let output = common::render_module("nix_shell").output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!("", actual);
+    Ok(())
+}
+
+#[test]
+fn invalid_env_variables() -> io::Result<()> {
+    let output = common::render_module("nix_shell")
+        .env("IN_NIX_SHELL", "something_wrong")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!("", actual);
+    Ok(())
+}
+
+#[test]
+fn pure_shell() -> io::Result<()> {
+    let output = common::render_module("nix_shell")
+        .env("IN_NIX_SHELL", "pure")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("via {} ", Color::Red.bold().paint("pure"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn impure_shell() -> io::Result<()> {
+    let output = common::render_module("nix_shell")
+        .env("IN_NIX_SHELL", "impure")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("via {} ", Color::Red.bold().paint("impure"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn lorri_shell() -> io::Result<()> {
+    let output = common::render_module("nix_shell")
+        .env("IN_NIX_SHELL", "1")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!("via {} ", Color::Red.bold().paint("impure"));
+    assert_eq!(expected, actual);
+    Ok(())
+}