feat: Enable the python module in virtual envs (#584)

This will enable the python module when a virtual environment has been
activated, this is detected via the `VIRTUAL_ENV` env var.

3 files changed, 21 insertions, 1 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index eab725b06..6a73d5b57 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -831,6 +831,7 @@ The module will be shown if any of the following conditions are met:
 - The current directory contains a file with the `.py` extension
 - The current directory contains a `Pipfile` file
 - The current directory contains a `tox.ini` file
+- A virtual environment is currently activated
 
 ### Options
 
diff --git a/src/modules/python.rs b/src/modules/python.rs
index 99ab78e46..491ad89a2 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -27,7 +27,9 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         .set_extensions(&["py"])
         .is_match();
 
-    if !is_py_project {
+    let is_venv = env::var("VIRTUAL_ENV").ok().is_some();
+
+    if !is_py_project && !is_venv {
         return None;
     }
 
diff --git a/tests/testsuite/python.rs b/tests/testsuite/python.rs
index d464fe91a..29a309b26 100644
--- a/tests/testsuite/python.rs
+++ b/tests/testsuite/python.rs
@@ -124,3 +124,20 @@ fn with_virtual_env() -> io::Result<()> {
     assert_eq!(expected, actual);
     Ok(())
 }
+
+#[test]
+#[ignore]
+fn with_active_venv() -> io::Result<()> {
+    let dir = tempfile::tempdir()?;
+
+    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.7.4 (my_venv)"));
+    assert_eq!(expected, actual);
+    Ok(())
+}