feat: Implement Python virtual environment display (#137)

docs/config/README.md

@@ -327,6 +327,8 @@ symbol = "🎁 "
 ## Python
 
 The `python` module shows the currently installed version of Python.
+It will also show the current Python virtual environment if one is
+activated.
 The module will be shown if any of the following conditions are met:
 
 - The current directory contains a `.python-version` file

src/modules/cmd_duration.rs

@@ -88,5 +88,4 @@ mod tests {
     fn test_1d() {
         assert_eq!(render_time(86400 as u64), "1d")
     }
-
 }

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")
+    }
 }

tests/testsuite/python.rs

@@ -1,7 +1,9 @@
-use ansi_term::Color;
+use std::env;
 use std::fs::File;
 use std::io;
 
+use ansi_term::Color;
+
 use crate::common;
 
 #[test]
@@ -71,3 +73,23 @@ fn folder_with_py_file() -> io::Result<()> {
     assert_eq!(expected, actual);
     Ok(())
 }
+
+#[test]
+#[ignore]
+fn with_virtual_env() -> io::Result<()> {
+    let dir = common::new_tempdir()?;
+    File::create(dir.path().join("main.py"))?;
+    let output = common::render_module("python")
+        .env("VIRTUAL_ENV", "/foo/bar/my_venv")
+        .arg("--path")
+        .arg(dir.path())
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!(
+        "via {} ",
+        Color::Yellow.bold().paint("🐍 v3.6.9(my_venv)")
+    );
+    assert_eq!(expected, actual);
+    Ok(())
+}