From d2946ddff7ca5ab289487dcaadb873add6b78d08 Mon Sep 17 00:00:00 2001
From: Asger Hautop Drewsen <asgerdrewsen@gmail.com>
Date: Sun, 14 Mar 2021 19:27:19 +0100
Subject: [PATCH] fix(python): Handle PyPy python version correctly (#2374)

* fix(python): Handle PyPy python version correctly

* refactor: rework Python version retrieval and formatting

Align Python version retrieval and formatting with established
Starship conventions.
---
 src/modules/python.rs | 51 ++++++++++++++++++++++++-------------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/src/modules/python.rs b/src/modules/python.rs
index 8fda37fc..e70dd02e 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -69,29 +69,30 @@ fn get_python_version(context: &Context, config: &PythonConfig) -> Option<String
     if config.pyenv_version_name {
         return Some(context.exec_cmd("pyenv", &["version-name"])?.stdout);
     };
-    let version = config.python_binary.0.iter().find_map(|binary| {
-        match context.exec_cmd(binary, &["--version"]) {
-            Some(output) => {
-                if output.stdout.is_empty() {
-                    Some(output.stderr)
-                } else {
-                    Some(output.stdout)
-                }
+    let version = config
+        .python_binary
+        .0
+        .iter()
+        .find_map(|binary| context.exec_cmd(binary, &["--version"]))
+        .map(|output| {
+            if output.stdout.is_empty() {
+                output.stderr
+            } else {
+                output.stdout
             }
-            None => None,
-        }
-    })?;
-    Some(format_python_version(&version))
+        })?;
+
+    format_python_version(&version)
 }
 
-fn format_python_version(python_stdout: &str) -> String {
-    format!(
-        "v{}",
-        python_stdout
-            .trim_start_matches("Python ")
-            .trim_end_matches(":: Anaconda, Inc.")
-            .trim()
-    )
+fn format_python_version(python_version: &str) -> Option<String> {
+    let version = python_version
+        // split into ["Python", "3.8.6", ...]
+        .split_whitespace()
+        // return "3.8.6"
+        .nth(1)?;
+
+    Some(format!("v{}", version))
 }
 
 fn get_python_virtual_env(context: &Context) -> Option<String> {
@@ -123,13 +124,19 @@ mod tests {
     #[test]
     fn test_format_python_version() {
         let input = "Python 3.7.2";
-        assert_eq!(format_python_version(input), "v3.7.2");
+        assert_eq!(format_python_version(input), Some("v3.7.2".to_string()));
     }
 
     #[test]
     fn test_format_python_version_anaconda() {
         let input = "Python 3.6.10 :: Anaconda, Inc.";
-        assert_eq!(format_python_version(input), "v3.6.10");
+        assert_eq!(format_python_version(input), Some("v3.6.10".to_string()));
+    }
+
+    #[test]
+    fn test_format_python_version_pypy() {
+        let input = "Python 3.7.9 (7e6e2bb30ac5fbdbd443619cae28c51d5c162a02, Nov 24 2020, 10:03:59)\n[PyPy 7.3.3-beta0 with GCC 10.2.0]";
+        assert_eq!(format_python_version(input), Some("v3.7.9".to_string()));
     }
 
     #[test]