feat: 解决python虚拟环境问题
This commit is contained in:
@@ -0,0 +1,386 @@
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
env,
|
||||
ffi::OsString,
|
||||
io,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use tokio::process::Command;
|
||||
|
||||
const VENV_DIRECTORY_NAMES: [&str; 3] = [".venv", "venv", "env"];
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreparedPython {
|
||||
pub interpreter: PathBuf,
|
||||
pub virtual_env: Option<PathBuf>,
|
||||
environment: BTreeMap<String, String>,
|
||||
remove_environment: Vec<&'static str>,
|
||||
}
|
||||
|
||||
impl PreparedPython {
|
||||
pub fn apply(&self, command: &mut Command) {
|
||||
for key in &self.remove_environment {
|
||||
command.env_remove(key);
|
||||
}
|
||||
command.envs(&self.environment);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn prepare_script_environment(
|
||||
working_dir: &str,
|
||||
interpreter: &str,
|
||||
configured_venv: Option<&str>,
|
||||
user_env: &BTreeMap<String, String>,
|
||||
) -> io::Result<PreparedPython> {
|
||||
if !is_python_interpreter(interpreter) {
|
||||
return Ok(PreparedPython {
|
||||
interpreter: resolve_interpreter(working_dir, interpreter),
|
||||
virtual_env: None,
|
||||
environment: user_env.clone(),
|
||||
remove_environment: Vec::new(),
|
||||
});
|
||||
}
|
||||
|
||||
let working_directory = absolute_path(Path::new(working_dir))?;
|
||||
let interpreter_path = resolve_interpreter(working_dir, interpreter);
|
||||
let explicit_venv = configured_venv
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(|value| resolve_against(&working_directory, value))
|
||||
.transpose()?;
|
||||
let detected_venv = explicit_venv
|
||||
.as_deref()
|
||||
.map(validate_venv)
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.or_else(|| detect_venv_from_interpreter(&interpreter_path))
|
||||
.or_else(|| {
|
||||
is_generic_python(interpreter)
|
||||
.then(|| detect_venv_from_ancestors(&working_directory))
|
||||
.flatten()
|
||||
});
|
||||
|
||||
let mut environment = BTreeMap::new();
|
||||
for (key, value) in user_env {
|
||||
if key.eq_ignore_ascii_case("PATH")
|
||||
|| key.eq_ignore_ascii_case("VIRTUAL_ENV")
|
||||
|| key.eq_ignore_ascii_case("VIRTUAL_ENV_PROMPT")
|
||||
|| (detected_venv.is_some() && key.eq_ignore_ascii_case("PYTHONHOME"))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
environment.insert(key.clone(), value.clone());
|
||||
}
|
||||
|
||||
if let Some(venv) = &detected_venv {
|
||||
let scripts_dir = venv.join("Scripts");
|
||||
let python = scripts_dir.join("python.exe");
|
||||
environment.insert(
|
||||
"VIRTUAL_ENV".to_owned(),
|
||||
venv.to_string_lossy().into_owned(),
|
||||
);
|
||||
environment.insert(
|
||||
"VIRTUAL_ENV_PROMPT".to_owned(),
|
||||
venv.file_name()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
);
|
||||
environment.insert(
|
||||
"PATH".to_owned(),
|
||||
build_venv_path(&scripts_dir, user_env)?
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
);
|
||||
environment.insert("PYTHONUNBUFFERED".to_owned(), "1".to_owned());
|
||||
return Ok(PreparedPython {
|
||||
interpreter: python,
|
||||
virtual_env: Some(venv.clone()),
|
||||
environment,
|
||||
remove_environment: vec![
|
||||
"PYTHONHOME",
|
||||
"PYTHONPATH",
|
||||
"VIRTUAL_ENV",
|
||||
"VIRTUAL_ENV_PROMPT",
|
||||
],
|
||||
});
|
||||
}
|
||||
|
||||
environment.insert("PYTHONUNBUFFERED".to_owned(), "1".to_owned());
|
||||
Ok(PreparedPython {
|
||||
interpreter: interpreter_path,
|
||||
virtual_env: None,
|
||||
environment,
|
||||
remove_environment: vec!["VIRTUAL_ENV", "VIRTUAL_ENV_PROMPT"],
|
||||
})
|
||||
}
|
||||
|
||||
pub fn validate_configured_venv(working_dir: &str, configured_venv: &str) -> io::Result<PathBuf> {
|
||||
let working_directory = absolute_path(Path::new(working_dir))?;
|
||||
let venv = resolve_against(&working_directory, configured_venv)?;
|
||||
validate_venv(&venv)?.ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("无效的 Python 虚拟环境: {}", venv.display()),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_venv(path: &Path) -> io::Result<Option<PathBuf>> {
|
||||
let root = absolute_path(path)?;
|
||||
let python = root.join("Scripts").join("python.exe");
|
||||
if !root.join("pyvenv.cfg").is_file() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("虚拟环境缺少 pyvenv.cfg: {}", root.display()),
|
||||
));
|
||||
}
|
||||
if !python.is_file() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("虚拟环境缺少 Scripts\\python.exe: {}", root.display()),
|
||||
));
|
||||
}
|
||||
Ok(Some(root))
|
||||
}
|
||||
|
||||
fn detect_venv_from_ancestors(start: &Path) -> Option<PathBuf> {
|
||||
for directory in start.ancestors() {
|
||||
for name in VENV_DIRECTORY_NAMES {
|
||||
let candidate = directory.join(name);
|
||||
if candidate.join("pyvenv.cfg").is_file()
|
||||
&& candidate.join("Scripts").join("python.exe").is_file()
|
||||
{
|
||||
return Some(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn detect_venv_from_interpreter(interpreter: &Path) -> Option<PathBuf> {
|
||||
let scripts = interpreter.parent()?;
|
||||
if !scripts
|
||||
.file_name()
|
||||
.is_some_and(|name| name.eq_ignore_ascii_case("Scripts"))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let root = scripts.parent()?;
|
||||
(root.join("pyvenv.cfg").is_file() && scripts.join("python.exe").is_file())
|
||||
.then(|| root.to_path_buf())
|
||||
}
|
||||
|
||||
fn build_venv_path(
|
||||
scripts_dir: &Path,
|
||||
user_env: &BTreeMap<String, String>,
|
||||
) -> io::Result<OsString> {
|
||||
let base = user_env
|
||||
.iter()
|
||||
.find(|(key, _)| key.eq_ignore_ascii_case("PATH"))
|
||||
.map(|(_, value)| OsString::from(value))
|
||||
.or_else(|| env::var_os("PATH"))
|
||||
.unwrap_or_default();
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
let mut paths = Vec::new();
|
||||
for path in std::iter::once(scripts_dir.to_path_buf()).chain(env::split_paths(&base)) {
|
||||
let key = path.to_string_lossy().to_ascii_lowercase();
|
||||
if seen.insert(key) {
|
||||
paths.push(path);
|
||||
}
|
||||
}
|
||||
env::join_paths(paths).map_err(io::Error::other)
|
||||
}
|
||||
|
||||
fn resolve_interpreter(working_dir: &str, interpreter: &str) -> PathBuf {
|
||||
let path = Path::new(interpreter);
|
||||
if path.is_absolute() || path.components().count() == 1 {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
Path::new(working_dir).join(path)
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_against(base: &Path, value: &str) -> io::Result<PathBuf> {
|
||||
let path = Path::new(value);
|
||||
let resolved = if path.is_absolute() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
base.join(path)
|
||||
};
|
||||
absolute_path(&resolved)
|
||||
}
|
||||
|
||||
fn absolute_path(path: &Path) -> io::Result<PathBuf> {
|
||||
std::path::absolute(path)
|
||||
}
|
||||
|
||||
fn is_python_interpreter(interpreter: &str) -> bool {
|
||||
let name = Path::new(interpreter)
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.unwrap_or(interpreter)
|
||||
.to_ascii_lowercase();
|
||||
name == "py.exe" || name == "py" || name.starts_with("python")
|
||||
}
|
||||
|
||||
fn is_generic_python(interpreter: &str) -> bool {
|
||||
Path::new(interpreter).components().count() == 1 && is_python_interpreter(interpreter)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn create_fake_venv(root: &Path) -> PathBuf {
|
||||
let venv = root.join(".venv");
|
||||
fs::create_dir_all(venv.join("Scripts")).unwrap();
|
||||
fs::write(venv.join("pyvenv.cfg"), "home = C:\\Python").unwrap();
|
||||
fs::write(venv.join("Scripts").join("python.exe"), "").unwrap();
|
||||
venv
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_venv_from_working_directory_and_builds_path() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let venv = create_fake_venv(temp.path());
|
||||
let working_dir = temp.path().join("project").join("src");
|
||||
fs::create_dir_all(&working_dir).unwrap();
|
||||
let mut user_env = BTreeMap::new();
|
||||
user_env.insert("CUSTOM_VALUE".to_owned(), "yes".to_owned());
|
||||
user_env.insert("PYTHONHOME".to_owned(), "C:\\wrong".to_owned());
|
||||
|
||||
let prepared = prepare_script_environment(
|
||||
working_dir.to_str().unwrap(),
|
||||
"python.exe",
|
||||
None,
|
||||
&user_env,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prepared.virtual_env, Some(venv.clone()));
|
||||
assert_eq!(
|
||||
prepared.interpreter,
|
||||
venv.join("Scripts").join("python.exe")
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.environment.get("VIRTUAL_ENV"),
|
||||
Some(&venv.to_string_lossy().into_owned())
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.environment.get("PYTHONUNBUFFERED"),
|
||||
Some(&"1".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_venv_is_validated() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let venv = create_fake_venv(temp.path());
|
||||
let prepared = prepare_script_environment(
|
||||
temp.path().to_str().unwrap(),
|
||||
"python.exe",
|
||||
Some(venv.to_str().unwrap()),
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prepared.virtual_env, Some(venv));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_explicit_venv_returns_error() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let missing = temp.path().join("missing-venv");
|
||||
let error = prepare_script_environment(
|
||||
temp.path().to_str().unwrap(),
|
||||
"python.exe",
|
||||
Some(missing.to_str().unwrap()),
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_python_interpreter_is_not_modified() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
create_fake_venv(temp.path());
|
||||
let prepared = prepare_script_environment(
|
||||
temp.path().to_str().unwrap(),
|
||||
"node.exe",
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prepared.virtual_env, None);
|
||||
assert!(prepared.environment.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn venv_environment_overrides_conflicting_python_variables() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let venv = create_fake_venv(temp.path());
|
||||
let mut user_env = BTreeMap::new();
|
||||
user_env.insert(
|
||||
"Path".to_owned(),
|
||||
temp.path()
|
||||
.join("custom-bin")
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
);
|
||||
user_env.insert("PYTHONHOME".to_owned(), "C:\\wrong".to_owned());
|
||||
user_env.insert("PYTHONUNBUFFERED".to_owned(), "0".to_owned());
|
||||
user_env.insert(
|
||||
"PYTHONPATH".to_owned(),
|
||||
"C:\\explicit-python-path".to_owned(),
|
||||
);
|
||||
|
||||
let prepared = prepare_script_environment(
|
||||
temp.path().to_str().unwrap(),
|
||||
"python.exe",
|
||||
None,
|
||||
&user_env,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(!prepared.environment.contains_key("PYTHONHOME"));
|
||||
assert_eq!(
|
||||
prepared.environment.get("PYTHONUNBUFFERED"),
|
||||
Some(&"1".to_owned())
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.environment.get("PYTHONPATH"),
|
||||
Some(&"C:\\explicit-python-path".to_owned())
|
||||
);
|
||||
let paths = env::split_paths(
|
||||
prepared
|
||||
.environment
|
||||
.get("PATH")
|
||||
.expect("venv PATH should be present"),
|
||||
)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(paths.first(), Some(&venv.join("Scripts")));
|
||||
assert!(paths.contains(&temp.path().join("custom-bin")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn global_python_removes_stale_virtual_env_markers() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let prepared = prepare_script_environment(
|
||||
temp.path().to_str().unwrap(),
|
||||
"python.exe",
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prepared.virtual_env, None);
|
||||
assert!(prepared.remove_environment.contains(&"VIRTUAL_ENV"));
|
||||
assert!(prepared.remove_environment.contains(&"VIRTUAL_ENV_PROMPT"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user