From 3669e389b6fe18b12a6ed2d0b3dcb1a9df9bc1f1 Mon Sep 17 00:00:00 2001
From: MaT1g3R <peijun.ma@protonmail.com>
Date: Sun, 11 Aug 2019 17:51:13 -0400
Subject: [PATCH] feat: Implement Python virtual environment display (#137)

---
 docs/config/README.md       |  2 ++
 src/modules/cmd_duration.rs |  1 -
 src/modules/python.rs       | 27 ++++++++++++++++++++++++++-
 tests/testsuite/python.rs   | 24 +++++++++++++++++++++++-
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index d86ae3391..004fe9521 100644
--- a/docs/config/README.md
+++ b/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
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")
+    }
 }
diff --git a/tests/testsuite/python.rs b/tests/testsuite/python.rs
index bc47c53a4..9be7c9c0e 100644
--- a/tests/testsuite/python.rs
+++ b/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(())
+}