feat(python): add option to replace generic venv-names with parent dir (#7112)

.github/config-schema.json

@@ -1436,7 +1436,8 @@
         "detect_folders": [],
         "detect_env_vars": [
           "VIRTUAL_ENV"
-        ]
+        ],
+        "generic_venv_names": []
       }
     },
     "quarto": {
@@ -5683,6 +5684,13 @@
           "default": [
             "VIRTUAL_ENV"
           ]
+        },
+        "generic_venv_names": {
+          "type": "array",
+          "items": {
+            "type": "string"
+          },
+          "default": []
         }
       },
       "additionalProperties": false

docs/config/README.md

@@ -4020,6 +4020,7 @@ By default, the module will be shown if any of the following conditions are met:
 | `detect_extensions`  | `['py', 'ipynb']`                                                                                            | 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                                              |
+| `generic_venv_names` | `[]`                                                                                                         | Which venv names should be replaced with the parent directory name.                   |
 | `disabled`           | `false`                                                                                                      | Disables the `python` module.                                                         |
 
 > [!TIP]
@@ -4039,13 +4040,13 @@ By default, the module will be shown if any of the following conditions are met:
 
 ### Variables
 
-| Variable     | Example         | Description                                |
-| ------------ | --------------- | ------------------------------------------ |
-| version      | `'v3.8.1'`      | The version of `python`                    |
-| symbol       | `'🐍 '`         | Mirrors the value of option `symbol`       |
-| style        | `'yellow bold'` | Mirrors the value of option `style`        |
-| pyenv_prefix | `'pyenv '`      | Mirrors the value of option `pyenv_prefix` |
-| virtualenv   | `'venv'`        | The current `virtualenv` name              |
+| Variable     | Example         | Description                                                                 |
+| ------------ | --------------- | --------------------------------------------------------------------------- |
+| version      | `'v3.8.1'`      | The version of `python`                                                     |
+| symbol       | `'🐍 '`         | Mirrors the value of option `symbol`                                        |
+| style        | `'yellow bold'` | Mirrors the value of option `style`                                         |
+| pyenv_prefix | `'pyenv '`      | Mirrors the value of option `pyenv_prefix`                                  |
+| virtualenv   | `'venv'`        | The current `virtualenv` name or the parent if matches `generic_venv_names` |
 
 ### Example
 

src/configs/python.rs

@@ -22,6 +22,7 @@ pub struct PythonConfig<'a> {
     pub detect_files: Vec<&'a str>,
     pub detect_folders: Vec<&'a str>,
     pub detect_env_vars: Vec<&'a str>,
+    pub generic_venv_names: Vec<&'a str>,
 }
 
 impl Default for PythonConfig<'_> {
@@ -51,6 +52,7 @@ impl Default for PythonConfig<'_> {
             ],
             detect_folders: vec![],
             detect_env_vars: vec!["VIRTUAL_ENV"],
+            generic_venv_names: vec![],
         }
     }
 }

src/modules/python.rs

@@ -56,7 +56,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
                     .map(Ok)
                 }
                 "virtualenv" => {
-                    let virtual_env = get_python_virtual_env(context);
+                    let virtual_env = get_python_virtual_env(context, &config);
                     virtual_env.as_ref().map(|e| Ok(e.trim().to_string()))
                 }
                 "pyenv_prefix" => Some(Ok(pyenv_prefix.to_string())),
@@ -123,13 +123,17 @@ fn parse_python_version(python_version_string: &str) -> Option<String> {
     Some(version.to_string())
 }
 
-fn get_python_virtual_env(context: &Context) -> Option<String> {
+fn get_python_virtual_env(context: &Context, config: &PythonConfig) -> Option<String> {
     context.get_env("VIRTUAL_ENV").and_then(|venv| {
-        get_prompt_from_venv(Path::new(&venv)).or_else(|| {
-            Path::new(&venv)
-                .file_name()
-                .map(|filename| String::from(filename.to_str().unwrap_or("")))
-        })
+        get_prompt_from_venv(Path::new(&venv))
+            .or_else(|| get_venv_from_path(Path::new(&venv)))
+            .and_then(|venv_name| {
+                if config.generic_venv_names.contains(&venv_name.as_str()) {
+                    get_venv_from_path(Path::new(&venv).parent()?)
+                } else {
+                    Some(venv_name)
+                }
+            })
     })
 }
 
@@ -141,6 +145,10 @@ fn get_prompt_from_venv(venv_path: &Path) -> Option<String> {
         .map(|prompt| String::from(prompt.trim_matches(&['(', ')'] as &[_])))
 }
 
+fn get_venv_from_path(venv_path: &Path) -> Option<String> {
+    venv_path.file_name()?.to_str().map(|s| s.to_string())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -394,6 +402,28 @@ Python 3.7.9 (7e6e2bb30ac5fbdbd443619cae28c51d5c162a02, Nov 24 2020, 10:03:59)
         dir.close()
     }
 
+    #[test]
+    fn with_active_venv_and_generic_venv_names() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("python")
+            .path(dir.path())
+            .env("VIRTUAL_ENV", "/foo/bar/my_venv")
+            .config(toml::toml! {
+                [python]
+                generic_venv_names = ["my_venv"]
+            })
+            .collect();
+
+        let expected = Some(format!(
+            "via {}",
+            Color::Yellow.bold().paint("🐍 v3.8.0 (bar) ")
+        ));
+
+        assert_eq!(actual, expected);
+        dir.close()
+    }
+
     #[test]
     fn with_different_env_var() -> io::Result<()> {
         let dir = tempfile::tempdir()?;
@@ -511,6 +541,40 @@ prompt = foo\nbar
         dir.close()
     }
 
+    #[test]
+    fn with_active_venv_and_generic_prompt() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        create_dir_all(dir.path().join("bar").join("my_venv"))?;
+        let mut venv_cfg = File::create(dir.path().join("bar").join("my_venv").join("pyvenv.cfg"))?;
+        venv_cfg.write_all(
+            br"
+home = something
+prompt = 'foo'
+        ",
+        )?;
+        venv_cfg.sync_all()?;
+
+        let actual = ModuleRenderer::new("python")
+            .path(dir.path())
+            .env(
+                "VIRTUAL_ENV",
+                dir.path().join("bar").join("my_venv").to_str().unwrap(),
+            )
+            .config(toml::toml! {
+                [python]
+                generic_venv_names = ["foo"]
+            })
+            .collect();
+
+        let expected = Some(format!(
+            "via {}",
+            Color::Yellow.bold().paint("🐍 v3.8.0 (bar) ")
+        ));
+
+        assert_eq!(actual, expected);
+        dir.close()
+    }
+
     fn check_python2_renders(dir: &tempfile::TempDir, starship_config: Option<toml::Table>) {
         let config = starship_config.unwrap_or(toml::toml! {
             [python]