feat: Implement Python virtual environment display (#137)

2 files changed, 26 insertions, 2 deletions

diff --git a/src/modules/cmd_duration.rs b/src/modules/cmd_duration.rs
index c38a9b601..2f9e0c32e 100644
--- a/src/modules/cmd_duration.rs
+++ b/src/modules/cmd_duration.rs
@@ -88,5 +88,4 @@ mod tests {
     fn test_1d() {
         assert_eq!(render_time(86400 as u64), "1d")
     }
-
 }
diff --git a/src/modules/python.rs b/src/modules/python.rs
index fe643d4c2..0e7b3c250 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -1,6 +1,9 @@
-use ansi_term::Color;
+use std::env;
+use std::path::Path;
 use std::process::Command;
 
+use ansi_term::Color;
+
 use super::{Context, Module};
 
 /// Creates a module with the current Python version
@@ -32,6 +35,8 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
             let formatted_version = format_python_version(&python_version);
             module.new_segment("symbol", PYTHON_CHAR);
             module.new_segment("version", &formatted_version);
+            get_python_virtual_env()
+                .map(|virtual_env| module.new_segment("virtualenv", &format!("({})", virtual_env)));
 
             Some(module)
         }
@@ -61,6 +66,14 @@ fn format_python_version(python_stdout: &str) -> String {
     format!("v{}", python_stdout.trim_start_matches("Python ").trim())
 }
 
+fn get_python_virtual_env() -> Option<String> {
+    env::var("VIRTUAL_ENV").ok().and_then(|venv| {
+        Path::new(&venv)
+            .file_name()
+            .map(|filename| String::from(filename.to_str().unwrap_or("")))
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -70,4 +83,16 @@ mod tests {
         let input = "Python 3.7.2";
         assert_eq!(format_python_version(input), "v3.7.2");
     }
+
+    #[test]
+    fn test_no_virtual_env() {
+        env::set_var("VIRTUAL_ENV", "");
+        assert_eq!(get_python_virtual_env(), None)
+    }
+
+    #[test]
+    fn test_virtual_env() {
+        env::set_var("VIRTUAL_ENV", "/foo/bar/my_venv");
+        assert_eq!(get_python_virtual_env().unwrap(), "my_venv")
+    }
 }