diff --git a/crates/pet-core/src/pyvenv_cfg.rs b/crates/pet-core/src/pyvenv_cfg.rs index 01f06015..4c37f6e2 100644 --- a/crates/pet-core/src/pyvenv_cfg.rs +++ b/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 { if let Some(ref file) = find(path) { @@ -77,16 +83,36 @@ fn parse(file: &Path) -> Option { 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 { + 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 +} diff --git a/crates/pet-python-utils/src/headers.rs b/crates/pet-python-utils/src/headers.rs index ebdd4fc4..e8848dc3 100644 --- a/crates/pet-python-utils/src/headers.rs +++ b/crates/pet-python-utils/src/headers.rs @@ -18,7 +18,12 @@ pub struct Headers { impl Headers { pub fn get_version(path: &Path) -> Option { - 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 { - 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 { // 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 { } 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 { + 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 = version + .splitn(3, ".") + .take(2) + .flat_map(str::parse::) + .collect(); + if parts.len() == 2 && (parts[0], parts[1]) == pyver { + return Some(version.to_string()); } + } else { + return Some(version.to_string()); } } } diff --git a/crates/pet-python-utils/src/version.rs b/crates/pet-python-utils/src/version.rs index b8b3ae4d..85db00eb 100644 --- a/crates/pet-python-utils/src/version.rs +++ b/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 { } 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.