This commit is contained in:
Zanie Blue 2025-06-11 09:59:38 -05:00
parent e0399dafd8
commit 7438c18662
3 changed files with 105 additions and 326 deletions

View file

@ -28,8 +28,7 @@ use crate::interpreter::{StatusCodeError, UnexpectedResponseError};
use crate::managed::ManagedPythonInstallations;
#[cfg(windows)]
use crate::microsoft_store::find_microsoft_store_pythons;
#[cfg(windows)]
use crate::pe_version::extract_version_from_pe;
use crate::pe_version::try_extract_version_from_pe;
use crate::virtualenv::Error as VirtualEnvError;
use crate::virtualenv::{
CondaEnvironmentKind, conda_environment_from_env, virtualenv_from_env,
@ -556,24 +555,25 @@ fn python_executables_from_search_path<'a>(
.into_iter()
.flatten(),
)
.filter(|path| {
// On Windows, try PE version extraction first as an optimization
// Skip this executable if PE version doesn't match
if let Some(pe_matches) = try_pe_version_check(path, Some(version)) {
if !pe_matches {
debug!(
"Skipping `{}` based on PE version mismatch",
path.display()
);
return false;
}
}
true
})
})
.into_iter()
.flatten()
})
.filter(move |path| {
// On Windows, skip executables with version in the PE metadata that
// does not match
if matches!(version, VersionRequest::Any | VersionRequest::Default) {
return true;
}
let Some(pe_python_version) = try_extract_version_from_pe(path) else {
return true;
};
if version.matches_version(&pe_python_version) {
return true;
}
debug!("Skipping interpreter at `{}`: PE version `{pe_python_version}` does not match `{version}`", path.display());
false
})
}
/// Find all acceptable `python3.x` minor versions.
@ -684,78 +684,23 @@ fn python_interpreters<'a>(
})
}
/// On Windows, try to extract version from PE metadata to potentially skip expensive interpreter queries.
/// Returns `None` if PE extraction fails or is not available, indicating fallback to full query.
#[cfg(windows)]
fn try_pe_version_check(path: &Path, request: Option<&VersionRequest>) -> Option<bool> {
let Some(request) = request else {
// No specific version requested, can't optimize
return None;
};
let pe_version = match extract_version_from_pe(path) {
Ok(Some(version)) => version,
Ok(None) => {
debug!("No version info found in PE file: {}", path.display());
return None;
}
Err(err) => {
debug!(
"Failed to extract PE version from {}: {}",
path.display(),
err
);
return None;
}
};
let matches = request.matches_version(&pe_version);
if matches {
debug!(
"PE version {} from {} matches request {}",
pe_version,
path.display(),
request
);
} else {
debug!(
"PE version {} from {} does not match request {}",
pe_version,
path.display(),
request
);
}
Some(matches)
}
/// Non-Windows version always returns None (no optimization available)
#[cfg(not(windows))]
fn try_pe_version_check(_path: &Path, _request: Option<&VersionRequest>) -> Option<bool> {
None
}
/// Lazily convert Python executables into interpreters.
fn python_interpreters_from_executables<'a>(
executables: impl Iterator<Item = Result<(PythonSource, PathBuf), Error>> + 'a,
cache: &'a Cache,
) -> impl Iterator<Item = Result<(PythonSource, Interpreter), Error>> + 'a {
executables.map(move |result| match result {
Ok((source, path)) => {
// Proceed with full interpreter query
Interpreter::query(&path, cache)
.map(|interpreter| (source, interpreter))
.inspect(|(source, interpreter)| {
debug!(
"Found `{}` at `{}` ({source})",
interpreter.key(),
path.display()
);
})
.map_err(|err| Error::Query(Box::new(err), path, source))
.inspect_err(|err| debug!("{err}"))
}
Ok((source, path)) => Interpreter::query(&path, cache)
.map(|interpreter| (source, interpreter))
.inspect(|(source, interpreter)| {
debug!(
"Found `{}` at `{}` ({source})",
interpreter.key(),
path.display()
);
})
.map_err(|err| Error::Query(Box::new(err), path, source))
.inspect_err(|err| debug!("{err}")),
Err(err) => Err(err),
})
}

View file

@ -13,7 +13,6 @@ pub use crate::environment::{InvalidEnvironmentKind, PythonEnvironment};
pub use crate::implementation::ImplementationName;
pub use crate::installation::{PythonInstallation, PythonInstallationKey};
pub use crate::interpreter::{BrokenSymlink, Error as InterpreterError, Interpreter};
pub use crate::pe_version::extract_version_from_pe;
pub use crate::pointer_size::PointerSize;
pub use crate::prefix::Prefix;
pub use crate::python_version::PythonVersion;

View file

@ -1,275 +1,110 @@
//! Extract version information from Windows PE executables.
//!
//! This module provides functionality to extract Python version information
//! directly from Windows PE executable files using the pelite crate, which
//! can be faster than executing the Python interpreter to query its version.
use std::path::Path;
#[cfg(target_os = "windows")]
use std::str::FromStr;
#[cfg(windows)]
use tracing::debug;
#[cfg(target_os = "windows")]
use thiserror::Error;
#[cfg(target_os = "windows")]
use pelite::{pe32::Pe as Pe32, pe64::Pe as Pe64};
use tracing::{debug, trace};
#[cfg(target_os = "windows")]
use crate::PythonVersion;
/// Extract Python version information from a Windows PE executable.
///
/// This function reads the PE file's version resource to extract version
/// information without executing the Python interpreter. This can be
/// significantly faster for version discovery.
///
/// # Arguments
///
/// * `path` - Path to the Python executable
///
/// # Returns
///
/// Returns `Ok(Some(PythonVersion))` if version information was successfully
/// extracted, `Ok(None)` if no version information was found, or an error
/// if the file could not be read or parsed.
#[cfg(target_os = "windows")]
pub fn extract_version_from_pe(path: &Path) -> Result<Option<PythonVersion>, std::io::Error> {
use pelite::FileMap;
#[derive(Debug, Error)]
pub(crate) enum PeVersionError {
#[error("Failed to read PE file: {0}")]
Io(#[from] std::io::Error),
debug!("Extracting version info from PE file: {}", path.display());
#[error("Failed to parse PE file: {0}")]
InvalidPeFormat(#[from] pelite::Error),
#[error("No version info found in PE file")]
NoVersionInfo,
#[error("Failed to parse version: {0}")]
InvalidVersion(String),
#[error("Version component too large: {0}")]
VersionComponentTooLarge(#[from] std::num::TryFromIntError),
}
/// Try to extract Python version information from a Windows PE executable.
///
/// On error, return [`None`].
/// On non-Windows platforms, this function always returns [`None`].
#[cfg(not(target_os = "windows"))]
pub(crate) fn try_extract_version_from_pe(_path: &Path) -> Option<crate::PythonVersion> {
None
}
/// Try to extract Python version information from a Windows PE executable.
///
/// On error, return [`None`].
#[cfg(target_os = "windows")]
pub(crate) fn try_extract_version_from_pe(path: &Path) -> Option<PythonVersion> {
match extract_version_from_pe(path) {
Ok(version) => Some(version),
Err(err) => {
debug!(
"Failed to extract version from PE file `{}`: {}",
path.display(),
err
);
None
}
}
}
/// Extract Python version information from a Windows PE executable.
#[cfg(target_os = "windows")]
fn extract_version_from_pe(path: &Path) -> Result<PythonVersion, PeVersionError> {
use pelite::FileMap;
trace!("Extracting version info from PE file: {}", path.display());
// Read the PE file
let map = FileMap::open(path)?;
// Parse as PE64 first, fall back to PE32 if needed
match parse_pe64_version(&map) {
Ok(version) => Ok(version),
Err(_) => parse_pe32_version(&map),
}
}
#[cfg(target_os = "windows")]
fn parse_pe64_version(map: &pelite::FileMap) -> Result<Option<PythonVersion>, std::io::Error> {
use pelite::pe64::PeFile;
let pe = PeFile::from_bytes(map).map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("Failed to parse PE64 file: {e}"),
)
})?;
extract_version_from_pe64_file(&pe)
}
#[cfg(target_os = "windows")]
fn parse_pe32_version(map: &pelite::FileMap) -> Result<Option<PythonVersion>, std::io::Error> {
use pelite::pe32::PeFile;
let pe = PeFile::from_bytes(map).map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("Failed to parse PE32 file: {e}"),
)
})?;
extract_version_from_pe32_file(&pe)
}
#[cfg(target_os = "windows")]
fn extract_version_from_pe64_file(
pe: &pelite::pe64::PeFile,
) -> Result<Option<PythonVersion>, std::io::Error> {
// Get resources from the PE file
let resources = pe.resources().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("Failed to read PE resources: {e}"),
)
})?;
// Parse as PE64 first, then fall back to PE32
let resources = read_pe64_resources(&map).or_else(|_| read_pe32_resources(&map))?;
// Try to get version info
let Ok(version_info) = resources.version_info() else {
debug!("No version info found in PE file");
return Ok(None);
};
let version_info = resources
.version_info()
.map_err(|_| PeVersionError::NoVersionInfo)?;
// Get the fixed file info which contains version numbers
let Some(fixed_info) = version_info.fixed() else {
debug!("No fixed version info found in PE file");
return Ok(None);
};
let fixed_info = version_info.fixed().ok_or(PeVersionError::NoVersionInfo)?;
// Extract version from the file version field
let file_version = fixed_info.dwFileVersion;
#[allow(clippy::cast_possible_truncation)]
let major = file_version.Major as u8;
#[allow(clippy::cast_possible_truncation)]
let minor = file_version.Minor as u8;
#[allow(clippy::cast_possible_truncation)]
let patch = file_version.Patch as u8;
let major = u8::try_from(file_version.Major)?;
let minor = u8::try_from(file_version.Minor)?;
let patch = u8::try_from(file_version.Patch)?;
// Validate that this looks like a Python version
if major == 0 || major > 10 || minor > 50 {
debug!(
"Version {}.{}.{} doesn't look like a Python version",
major, minor, patch
);
return Ok(None);
}
debug!("Extracted Python version: {}.{}.{}", major, minor, patch);
match PythonVersion::from_str(&format!("{major}.{minor}.{patch}")) {
Ok(version) => Ok(Some(version)),
Err(e) => {
debug!(
"Failed to parse version {}.{}.{}: {}",
major, minor, patch, e
);
Ok(None)
}
}
PythonVersion::from_str(&format!("{major}.{minor}.{patch}"))
.map_err(PeVersionError::InvalidVersion)
}
#[cfg(target_os = "windows")]
fn extract_version_from_pe32_file(
pe: &pelite::pe32::PeFile,
) -> Result<Option<PythonVersion>, std::io::Error> {
// Get resources from the PE file
let resources = pe.resources().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("Failed to read PE resources: {e}"),
)
})?;
fn read_pe64_resources(
map: &pelite::FileMap,
) -> Result<pelite::resources::Resources, PeVersionError> {
use pelite::pe64::{Pe, PeFile};
// Try to get version info
let Ok(version_info) = resources.version_info() else {
debug!("No version info found in PE file");
return Ok(None);
};
// Get the fixed file info which contains version numbers
let Some(fixed_info) = version_info.fixed() else {
debug!("No fixed version info found in PE file");
return Ok(None);
};
// Extract version from the file version field
let file_version = fixed_info.dwFileVersion;
#[allow(clippy::cast_possible_truncation)]
let major = file_version.Major as u8;
#[allow(clippy::cast_possible_truncation)]
let minor = file_version.Minor as u8;
#[allow(clippy::cast_possible_truncation)]
let patch = file_version.Patch as u8;
// Validate that this looks like a Python version
if major == 0 || major > 10 || minor > 50 {
debug!(
"Version {}.{}.{} doesn't look like a Python version",
major, minor, patch
);
return Ok(None);
}
debug!("Extracted Python version: {}.{}.{}", major, minor, patch);
match PythonVersion::from_str(&format!("{major}.{minor}.{patch}")) {
Ok(version) => Ok(Some(version)),
Err(e) => {
debug!(
"Failed to parse version {}.{}.{}: {}",
major, minor, patch, e
);
Ok(None)
}
}
let pe = PeFile::from_bytes(map)?;
Ok(pe.resources()?)
}
/// Extract version information from a Windows PE executable.
///
/// On non-Windows platforms, this function always returns `Ok(None)`.
#[cfg(not(target_os = "windows"))]
pub fn extract_version_from_pe(_path: &Path) -> Result<Option<PythonVersion>, std::io::Error> {
Ok(None)
}
#[test]
fn test_basic_pe_version_functionality() {
use std::str::FromStr;
// Basic test for the non-Windows version
#[cfg(not(target_os = "windows"))]
{
let result = extract_version_from_pe(Path::new("test.exe"));
assert_eq!(result.unwrap(), None);
}
// Test PythonVersion parsing
assert!(PythonVersion::from_str("3.12.0").is_ok());
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
#[cfg(target_os = "windows")]
use std::io::Write;
#[cfg(target_os = "windows")]
use tempfile::NamedTempFile;
#[test]
#[cfg(target_os = "windows")]
fn test_extract_version_from_nonexistent_file() {
let result = extract_version_from_pe(Path::new("nonexistent.exe"));
assert!(result.is_err());
}
#[test]
#[cfg(target_os = "windows")]
fn test_extract_version_from_invalid_pe_file() {
// Create a temporary file with invalid PE content
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"Not a PE file").unwrap();
temp_file.flush().unwrap();
let result = extract_version_from_pe(temp_file.path());
assert!(result.is_err());
}
#[test]
#[cfg(not(target_os = "windows"))]
fn test_extract_version_non_windows() {
let result = extract_version_from_pe(Path::new("python.exe"));
assert_eq!(result.unwrap(), None);
}
#[test]
fn test_version_validation() {
// Test that valid Python versions work
assert!(PythonVersion::from_str("3.12.0").is_ok());
assert!(PythonVersion::from_str("3.9.18").is_ok());
assert!(PythonVersion::from_str("3.13.1").is_ok());
// Test some edge cases that should still work
assert!(PythonVersion::from_str("0.1.0").is_ok()); // PythonVersion allows this
// Test malformed versions
assert!(PythonVersion::from_str("not.a.version").is_err());
assert!(PythonVersion::from_str("").is_err());
}
#[test]
fn test_always_runs() {
// This test should always run regardless of platform
// Test that the non-Windows version works
let result = extract_version_from_pe(Path::new("fake.exe"));
#[cfg(not(target_os = "windows"))]
assert_eq!(result.unwrap(), None);
#[cfg(target_os = "windows")]
assert!(result.is_err() || result.unwrap().is_none());
}
#[cfg(target_os = "windows")]
fn read_pe32_resources(
map: &pelite::FileMap,
) -> Result<pelite::resources::Resources, PeVersionError> {
use pelite::pe32::{Pe, PeFile};
let pe = PeFile::from_bytes(map)?;
Ok(pe.resources()?)
}