feat(python): Smarter python binary usage (#1947)

Update the python module to try multiple python binaries when
determining the version. With the new logic if starship doesn't find
`python` on the `PATH`, which is the default for some Linux Distros, it
will fallback to `python3` and then `python2`.

4 files changed, 79 insertions, 25 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index 6d6547d89..d5d8b44f1 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -2009,16 +2009,34 @@ The module will be shown if any of the following conditions are met:
 
 ### Options
 
-| Option               | Default                                                                 | Description                                                                   |
-| -------------------- | ----------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
-| `format`             | `'via [${symbol}${pyenv_prefix}${version}( \($virtualenv\))]($style) '` | The format for the module.                                                    |
-| `symbol`             | `"🐍 "`                                                                 | A format string representing the symbol of Python                             |
-| `style`              | `"yellow bold"`                                                         | The style for the module.                                                     |
-| `pyenv_version_name` | `false`                                                                 | Use pyenv to get Python version                                               |
-| `pyenv_prefix`       | `pyenv `                                                                | Prefix before pyenv version display, only used if pyenv is used               |
-| `scan_for_pyfiles`   | `true`                                                                  | If false, Python files in the current directory will not show this module.    |
-| `python_binary`      | `python`                                                                | Configures the python binary that Starship executes when getting the version. |
-| `disabled`           | `false`                                                                 | Disables the `python` module.                                                 |
+| Option               | Default                                                                 | Description                                                                            |
+| -------------------- | ----------------------------------------------------------------------- | -------------------------------------------------------------------------------------- |
+| `format`             | `'via [${symbol}${pyenv_prefix}${version}( \($virtualenv\))]($style) '` | The format for the module.                                                             |
+| `symbol`             | `"🐍 "`                                                                 | A format string representing the symbol of Python                                      |
+| `style`              | `"yellow bold"`                                                         | The style for the module.                                                              |
+| `pyenv_version_name` | `false`                                                                 | Use pyenv to get Python version                                                        |
+| `pyenv_prefix`       | `pyenv `                                                                | Prefix before pyenv version display, only used if pyenv is used                        |
+| `scan_for_pyfiles`   | `true`                                                                  | If false, Python files in the current directory will not show this module.             |
+| `python_binary`      | `["python", "python3, "python2"]`                                       | Configures the python binaries that Starship should executes when getting the version. |
+| `disabled`           | `false`                                                                 | Disables the `python` module.                                                          |
+
+::: tip
+
+The `python_binary` variable accepts either a string or a list of strings.
+Starship will try executing each binary until it gets a result. Note you can
+only change the binary that Starship executes to get the version of Python not
+the arguments that are used.
+
+The default values and order for `python_binary` was chosen to first identify
+the Python version in a virtualenv/conda environments (which currently still
+add a `python`, no matter if it points to `python3` or `python2`). This has the
+side effect that if you still have a system Python 2 installed, it may be
+picked up before any Python 3 (at least on Linux Distros that always symlink
+`/usr/bin/python` to Python 2). If you do not work with Python 2 anymore but
+cannot remove the system Python 2, changing this to `"python3"` will hide any
+Python version 2, see example below.
+
+:::
 
 ### Variables
 
@@ -2041,15 +2059,11 @@ symbol = "👾 "
 pyenv_version_name = true
 ```
 
-Using the `python3` binary to get the version.
-
-Note - The `python_binary` variable changes the binary that Starship executes
-to get the version of Python, it doesn't change the arguments that are used.
-
 ```toml
 # ~/.config/starship.toml
 
 [python]
+# Only use the `python3` binary to get the version.
 python_binary = "python3"
 ```
 
diff --git a/src/configs/python.rs b/src/configs/python.rs
index ecad467c9..aec645daf 100644
--- a/src/configs/python.rs
+++ b/src/configs/python.rs
@@ -1,4 +1,4 @@
-use crate::config::{ModuleConfig, RootModuleConfig};
+use crate::config::{ModuleConfig, RootModuleConfig, VecOr};
 
 use starship_module_config_derive::ModuleConfig;
 
@@ -6,7 +6,7 @@ use starship_module_config_derive::ModuleConfig;
 pub struct PythonConfig<'a> {
     pub pyenv_version_name: bool,
     pub pyenv_prefix: &'a str,
-    pub python_binary: &'a str,
+    pub python_binary: VecOr<&'a str>,
     pub scan_for_pyfiles: bool,
     pub format: &'a str,
     pub style: &'a str,
@@ -19,7 +19,7 @@ impl<'a> RootModuleConfig<'a> for PythonConfig<'a> {
         PythonConfig {
             pyenv_version_name: false,
             pyenv_prefix: "pyenv ",
-            python_binary: "python",
+            python_binary: VecOr(vec!["python", "python3", "python2"]),
             scan_for_pyfiles: true,
             format: "via [${symbol}${pyenv_prefix}${version}( \\($virtualenv\\))]($style) ",
             style: "yellow bold",
diff --git a/src/modules/python.rs b/src/modules/python.rs
index b8affc7df..0d18235a5 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -45,7 +45,11 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let python_version = if config.pyenv_version_name {
         utils::exec_cmd("pyenv", &["version-name"])?.stdout
     } else {
-        let version = get_python_version(&config.python_binary)?;
+        let version = config
+            .python_binary
+            .0
+            .iter()
+            .find_map(|binary| get_python_version(binary))?;
         format_python_version(&version)
     };
     let virtual_env = get_python_virtual_env(context);
@@ -164,6 +168,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -175,6 +180,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -186,6 +192,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -197,6 +204,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -208,6 +216,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -219,6 +228,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -230,6 +240,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -241,6 +252,7 @@ mod tests {
         check_python2_renders(&dir, None);
         check_python3_renders(&dir, None);
         check_pyenv_renders(&dir, None);
+        check_multiple_binaries_renders(&dir, None);
         dir.close()
     }
 
@@ -269,6 +281,7 @@ mod tests {
 
         let config = toml::toml! {
             [python]
+            python_binary = "python2"
             scan_for_pyfiles = false
         };
 
@@ -289,6 +302,14 @@ mod tests {
             scan_for_pyfiles = false
         };
         check_pyenv_renders(&dir, Some(config_pyenv));
+
+        let config_multi = toml::toml! {
+            [python]
+            python_binary = ["python", "python3"]
+            scan_for_pyfiles = false
+        };
+        check_multiple_binaries_renders(&dir, Some(config_multi));
+
         dir.close()
     }
 
@@ -303,7 +324,7 @@ mod tests {
 
         let expected = Some(format!(
             "via {} ",
-            Color::Yellow.bold().paint("🐍 v2.7.17 (my_venv)")
+            Color::Yellow.bold().paint("🐍 v3.8.0 (my_venv)")
         ));
 
         assert_eq!(actual, expected);
@@ -321,7 +342,7 @@ mod tests {
 
         let expected = Some(format!(
             "via {} ",
-            Color::Yellow.bold().paint("🐍 v2.7.17 (my_venv)")
+            Color::Yellow.bold().paint("🐍 v3.8.0 (my_venv)")
         ));
 
         assert_eq!(actual, expected);
@@ -336,7 +357,7 @@ mod tests {
         venv_cfg.write_all(
             br#"
 home = something
-prompt = 'foo' 
+prompt = 'foo'
         "#,
         )?;
         venv_cfg.sync_all()?;
@@ -348,7 +369,7 @@ prompt = 'foo'
 
         let expected = Some(format!(
             "via {} ",
-            Color::Yellow.bold().paint("🐍 v2.7.17 (foo)")
+            Color::Yellow.bold().paint("🐍 v3.8.0 (foo)")
         ));
 
         assert_eq!(actual, expected);
@@ -358,7 +379,7 @@ prompt = 'foo'
     fn check_python2_renders(dir: &tempfile::TempDir, starship_config: Option<toml::Value>) {
         let config = starship_config.unwrap_or(toml::toml! {
             [python]
-            python_binary = "python"
+            python_binary = "python2"
         });
 
         let actual = ModuleRenderer::new("python")
@@ -385,6 +406,24 @@ prompt = 'foo'
         assert_eq!(expected, actual);
     }
 
+    fn check_multiple_binaries_renders(
+        dir: &tempfile::TempDir,
+        starship_config: Option<toml::Value>,
+    ) {
+        let config = starship_config.unwrap_or(toml::toml! {
+             [python]
+             python_binary = ["python", "python3"]
+        });
+
+        let actual = ModuleRenderer::new("python")
+            .path(dir.path())
+            .config(config)
+            .collect();
+
+        let expected = Some(format!("via {} ", Color::Yellow.bold().paint("🐍 v3.8.0")));
+        assert_eq!(expected, actual);
+    }
+
     fn check_pyenv_renders(dir: &tempfile::TempDir, starship_config: Option<toml::Value>) {
         let config = starship_config.unwrap_or(toml::toml! {
              [python]
diff --git a/src/utils.rs b/src/utils.rs
index 428d10982..00391cfc1 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -137,7 +137,8 @@ active boot switches: -d:release\n",
             stdout: String::from("system\n"),
             stderr: String::default(),
         }),
-        "python --version" => Some(CommandOutput {
+        "python --version" => None,
+        "python2 --version" => Some(CommandOutput {
             stdout: String::default(),
             stderr: String::from("Python 2.7.17\n"),
         }),