From 8a2d944bc80a0a8358c30e639e8af51c1d45d973 Mon Sep 17 00:00:00 2001
From: Thomas O'Donnell <andytom@users.noreply.github.com>
Date: Sun, 18 Aug 2024 18:11:13 +0200
Subject: [PATCH] feat(python): Add detect env vars option  (#4486)

feat(python): Configure different detect env vars

Allow the env vars that trigger the python module to be configured, if
an empty list is passed then the module will fall back to just
triggering based on the configured files, folders and extensions.
---
 .github/config-schema.json | 12 +++++++++++
 docs/config/README.md      |  1 +
 src/configs/python.rs      |  2 ++
 src/modules/python.rs      | 41 ++++++++++++++++++++++++++++++++++++--
 4 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/.github/config-schema.json b/.github/config-schema.json
index e3cb4b3a5..2733295f9 100644
--- a/.github/config-schema.json
+++ b/.github/config-schema.json
@@ -1396,6 +1396,9 @@
     },
     "python": {
       "default": {
+        "detect_env_vars": [
+          "VIRTUAL_ENV"
+        ],
         "detect_extensions": [
           "py"
         ],
@@ -5258,6 +5261,15 @@
           "items": {
             "type": "string"
           }
+        },
+        "detect_env_vars": {
+          "default": [
+            "VIRTUAL_ENV"
+          ],
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
         }
       },
       "additionalProperties": false
diff --git a/docs/config/README.md b/docs/config/README.md
index 7cc1008db..2246edcc5 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -3673,6 +3673,7 @@ By default, the module will be shown if any of the following conditions are met:
 | `detect_extensions`  | `['py']`                                                                                                     | Which extensions should trigger this module                                            |
 | `detect_files`       | `['.python-version', 'Pipfile', '__init__.py', 'pyproject.toml', 'requirements.txt', 'setup.py', 'tox.ini']` | Which filenames should trigger this module                                             |
 | `detect_folders`     | `[]`                                                                                                         | Which folders should trigger this module                                               |
+| `detect_env_vars`    | `["VIRTUAL_ENV"]`                                                                                            | Which environmental variables should trigger this module                               |
 | `disabled`           | `false`                                                                                                      | Disables the `python` module.                                                          |
 
 ::: tip
diff --git a/src/configs/python.rs b/src/configs/python.rs
index bcfa713df..3654fa82a 100644
--- a/src/configs/python.rs
+++ b/src/configs/python.rs
@@ -21,6 +21,7 @@ pub struct PythonConfig<'a> {
     pub detect_extensions: Vec<&'a str>,
     pub detect_files: Vec<&'a str>,
     pub detect_folders: Vec<&'a str>,
+    pub detect_env_vars: Vec<&'a str>,
 }
 
 impl<'a> Default for PythonConfig<'a> {
@@ -45,6 +46,7 @@ impl<'a> Default for PythonConfig<'a> {
                 "__init__.py",
             ],
             detect_folders: vec![],
+            detect_env_vars: vec!["VIRTUAL_ENV"],
         }
     }
 }
diff --git a/src/modules/python.rs b/src/modules/python.rs
index 299741638..d531362b1 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -19,9 +19,10 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         .set_folders(&config.detect_folders)
         .is_match();
 
-    let is_venv = context.get_env("VIRTUAL_ENV").is_some();
+    let has_env_vars =
+        !config.detect_env_vars.is_empty() && context.detect_env_vars(&config.detect_env_vars);
 
-    if !is_py_project && !is_venv {
+    if !is_py_project && !has_env_vars {
         return None;
     };
 
@@ -371,6 +372,42 @@ Python 3.7.9 (7e6e2bb30ac5fbdbd443619cae28c51d5c162a02, Nov 24 2020, 10:03:59)
         dir.close()
     }
 
+    #[test]
+    fn with_different_env_var() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("python")
+            .path(dir.path())
+            .env("MY_ENV_VAR", "my_env_var")
+            .config(toml::toml! {
+                [python]
+                detect_env_vars = ["MY_ENV_VAR"]
+            })
+            .collect();
+
+        let expected = Some(format!("via {}", Color::Yellow.bold().paint("🐍 v3.8.0 ")));
+
+        assert_eq!(actual, expected);
+        dir.close()
+    }
+
+    #[test]
+    fn with_no_env_var() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("python")
+            .path(dir.path())
+            .env("VIRTUAL_ENV", "env_var")
+            .config(toml::toml! {
+                [python]
+                detect_env_vars = []
+            })
+            .collect();
+
+        assert_eq!(actual, None);
+        dir.close()
+    }
+
     #[test]
     fn with_active_venv_and_prompt() -> io::Result<()> {
         let dir = tempfile::tempdir()?;