Fix env version detection on systems with multiple system Pythons

crates/pet-core/src/pyvenv_cfg.rs

@@ -20,11 +20,17 @@ const PYVENV_CONFIG_FILE: &str = "pyvenv.cfg";
 #[derive(Debug)]
 pub struct PyVenvCfg {
     pub version: String,
+    pub version_major: u64,
+    pub version_minor: u64,
 }
 
 impl PyVenvCfg {
-    fn new(version: String) -> Self {
-        Self { version }
+    fn new(version: String, version_major: u64, version_minor: u64) -> Self {
+        Self {
+            version,
+            version_major,
+            version_minor,
+        }
     }
     pub fn find(path: &Path) -> Option<Self> {
         if let Some(ref file) = find(path) {
@@ -77,16 +83,36 @@ fn parse(file: &Path) -> Option<PyVenvCfg> {
         if !line.contains("version") {
             continue;
         }
-        if let Some(captures) = VERSION.captures(line) {
-            if let Some(value) = captures.get(1) {
-                return Some(PyVenvCfg::new(value.as_str().to_string()));
-            }
+        if let Some(cfg) = parse_version(line, &VERSION) {
+            return Some(cfg);
         }
-        if let Some(captures) = VERSION_INFO.captures(line) {
-            if let Some(value) = captures.get(1) {
-                return Some(PyVenvCfg::new(value.as_str().to_string()));
-            }
+        if let Some(cfg) = parse_version(line, &VERSION_INFO) {
+            return Some(cfg);
         }
     }
     None
 }
+
+fn parse_version(line: &str, regex: &Regex) -> Option<PyVenvCfg> {
+    if let Some(captures) = regex.captures(line) {
+        if let Some(value) = captures.get(1) {
+            let version = value.as_str();
+            let parts: Vec<&str> = version.splitn(3, ".").take(2).collect();
+            // .expect() below is OK because the version regex
+            // guarantees there are at least two digits.
+            let version_major = parts[0]
+                .parse()
+                .expect("python major version to be an integer");
+            let version_minor = parts[1]
+                .parse()
+                .expect("python minor version to be an integer");
+            return Some(PyVenvCfg::new(
+                version.to_string(),
+                version_major,
+                version_minor,
+            ));
+        }
+    }
+
+    None
+}

crates/pet-python-utils/src/headers.rs

@@ -18,7 +18,12 @@ pub struct Headers {
 
 impl Headers {
     pub fn get_version(path: &Path) -> Option<String> {
-        get_version(path)
+        let mut path = path.to_path_buf();
+        let bin = if cfg!(windows) { "Scripts" } else { "bin" };
+        if path.ends_with(bin) {
+            path.pop();
+        }
+        get_version(&path, None)
     }
 }
 
@@ -28,23 +33,16 @@ impl Headers {
 // /* Version as a string */
 // #define PY_VERSION              "3.10.2"
 // /*--end constants--*/
-pub fn get_version(path: &Path) -> Option<String> {
-    let mut path = path.to_path_buf();
-    let bin = if cfg!(windows) { "Scripts" } else { "bin" };
-    if path.ends_with(bin) {
-        path.pop();
-    }
+pub fn get_version(sys_prefix: &Path, pyver: Option<(u64, u64)>) -> Option<String> {
     // Generally the files are in Headers in windows and include in unix
     // However they can also be in Headers on Mac (command line tools python, hence make no assumptions)
-    for headers_path in [path.join("Headers"), path.join("include")] {
-        let patchlevel_h = headers_path.join("patchlevel.h");
-        let mut contents = "".to_string();
-        if let Ok(result) = fs::read_to_string(patchlevel_h) {
-            contents = result;
-        } else if !headers_path.exists() {
-            // TODO: Remove this check, unnecessary, as we try to read the dir below.
-            // Such a path does not exist, get out.
+    for headers_path in [sys_prefix.join("Headers"), sys_prefix.join("include")] {
+        if !headers_path.exists() {
             continue;
+        }
+        let patchlevel_h = headers_path.join("patchlevel.h");
+        if let Some(version) = valid_version_from_header(&patchlevel_h, pyver) {
+            return Some(version);
         } else {
             // Try the other path
             // Sometimes we have it in a sub directory such as `python3.10` or `pypy3.9`
@@ -57,18 +55,32 @@ pub fn get_version(path: &Path) -> Option<String> {
                     }
                     let path = path.path();
                     let patchlevel_h = path.join("patchlevel.h");
-                    if let Ok(result) = fs::read_to_string(patchlevel_h) {
-                        contents = result;
-                        break;
+                    if let Some(version) = valid_version_from_header(&patchlevel_h, pyver) {
+                        return Some(version);
                     }
                 }
             }
         }
-        for line in contents.lines() {
-            if let Some(captures) = VERSION.captures(line) {
-                if let Some(value) = captures.get(1) {
-                    return Some(value.as_str().to_string());
+    }
+    None
+}
+
+fn valid_version_from_header(header: &Path, pyver: Option<(u64, u64)>) -> Option<String> {
+    let contents = fs::read_to_string(header).ok()?;
+    for line in contents.lines() {
+        if let Some(captures) = VERSION.captures(line) {
+            let version = captures.get(1)?.as_str();
+            if let Some(pyver) = pyver {
+                let parts: Vec<u64> = version
+                    .splitn(3, ".")
+                    .take(2)
+                    .flat_map(str::parse::<u64>)
+                    .collect();
+                if parts.len() == 2 && (parts[0], parts[1]) == pyver {
+                    return Some(version.to_string());
                 }
+            } else {
+                return Some(version.to_string());
             }
         }
     }

crates/pet-python-utils/src/version.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use crate::headers::Headers;
+use crate::headers::{self, Headers};
 use log::{trace, warn};
 use pet_core::pyvenv_cfg::PyVenvCfg;
 use pet_fs::path::resolve_symlink;
@@ -38,7 +38,13 @@ pub fn from_creator_for_virtual_env(prefix: &Path) -> Option<String> {
         } else {
             // Assume the python environment used to create this virtual env is a regular install of Python.
             // Try to get the version of that environment.
-            from_header_files(parent_dir)
+            let sys_root = parent_dir.parent()?;
+            let pyver = if let Some(pyvenvcfg) = PyVenvCfg::find(prefix) {
+                Some((pyvenvcfg.version_major, pyvenvcfg.version_minor))
+            } else {
+                None
+            };
+            headers::get_version(sys_root, pyver)
         }
     } else if cfg!(windows) {
         // Only on windows is it difficult to get the creator of the virtual environment.