detect nix-shell env

2 files changed, 16 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index b839325..1a77759 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -50,12 +50,13 @@ fn pista(zsh: bool) -> String {
         )
     } else {
         format!(
-            "{cwd} {branch} {status}\n{venv}{pchar} ",
+            "{cwd} {branch} {status}\n{venv}{nix}{pchar} ",
             cwd = cwd,
             branch = branch,
             status = status,
             venv = venv,
-            pchar = prompt_char
+            pchar = prompt_char,
+            nix = venv::in_nix_shell()
         )
     }
 }
diff --git a/src/venv.rs b/src/venv.rs
index 7906abc..da6830c 100644
--- a/src/venv.rs
+++ b/src/venv.rs
@@ -17,3 +17,16 @@ pub fn get_name() -> colored::ColoredString {
     return "".white()
 }
 
+pub fn in_nix_shell() -> colored::ColoredString {
+    match env::var("IN_NIX_SHELL") {
+        Ok(p) => {
+            if p == "pure" {
+                "(nix) ".green()
+            } else {
+                "(nix) ".red()
+            }
+        }
+        _ => return "".white(),
+    }
+}
+