From ac5a2b2122f5e092923e1529c6634917d0dce771 Mon Sep 17 00:00:00 2001 From: Elvis Pranskevichus Date: Tue, 1 Apr 2025 10:12:16 -0700 Subject: [PATCH] python: Add recognition of worktree -> venv links (#26759) Python envs created with virtualenvwrapper have the ability to link to the project directory (via the `.project` file in the virtualenv). `python-environment-tools` supports this and reports the project path. Additionally, some tools (e.g virtualfish) recognize special "environment activation" files in project directories (`.venv` by default) [1]. Use the above information to sort reported Python toolchains so that the correct env for a given worktree is at the top. [1] https://virtualfish.readthedocs.io/en/latest/plugins.html#auto-activation-auto-activation Release Notes: - python: Improved detection of virtualenvwrapper environments in work trees --- crates/languages/src/python.rs | 99 ++++++++++++++++++++++++++-------- 1 file changed, 76 insertions(+), 23 deletions(-) diff --git a/crates/languages/src/python.rs b/crates/languages/src/python.rs index 7b0400eb9c..9e8baf2470 100644 --- a/crates/languages/src/python.rs +++ b/crates/languages/src/python.rs @@ -30,6 +30,8 @@ use std::{ borrow::Cow, ffi::OsString, fmt::Write, + fs, + io::{self, BufRead}, path::{Path, PathBuf}, sync::Arc, }; @@ -635,6 +637,7 @@ static ENV_PRIORITY_LIST: &'static [PythonEnvironmentKind] = &[ PythonEnvironmentKind::VirtualEnvWrapper, PythonEnvironmentKind::Venv, PythonEnvironmentKind::VirtualEnv, + PythonEnvironmentKind::PyenvVirtualEnv, PythonEnvironmentKind::Pixi, PythonEnvironmentKind::Conda, PythonEnvironmentKind::Pyenv, @@ -654,6 +657,19 @@ fn env_priority(kind: Option) -> usize { } } +/// Return the name of environment declared in Option { + fs::File::open(worktree_root.join(".venv")) + .and_then(|file| { + let mut venv_name = String::new(); + io::BufReader::new(file).read_line(&mut venv_name)?; + Ok(venv_name.trim().to_string()) + }) + .ok() +} + #[async_trait] impl ToolchainLister for PythonToolchainProvider { async fn list( @@ -669,7 +685,7 @@ impl ToolchainLister for PythonToolchainProvider { &environment, ); let mut config = Configuration::default(); - config.workspace_directories = Some(vec![worktree_root]); + config.workspace_directories = Some(vec![worktree_root.clone()]); for locator in locators.iter() { locator.configure(&config); } @@ -683,29 +699,66 @@ impl ToolchainLister for PythonToolchainProvider { .ok() .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard)); + let wr = worktree_root; + let wr_venv = get_worktree_venv_declaration(&wr); + // Sort detected environments by: + // environment name matching activation file (/.venv) + // environment project dir matching worktree_root + // general env priority + // environment path matching the CONDA_PREFIX env var + // executable path toolchains.sort_by(|lhs, rhs| { - env_priority(lhs.kind) - .cmp(&env_priority(rhs.kind)) - .then_with(|| { - if lhs.kind == Some(PythonEnvironmentKind::Conda) { - environment - .get_env_var("CONDA_PREFIX".to_string()) - .map(|conda_prefix| { - let is_match = |exe: &Option| { - exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix)) - }; - match (is_match(&lhs.executable), is_match(&rhs.executable)) { - (true, false) => Ordering::Less, - (false, true) => Ordering::Greater, - _ => Ordering::Equal, - } - }) - .unwrap_or(Ordering::Equal) - } else { - Ordering::Equal - } - }) - .then_with(|| lhs.executable.cmp(&rhs.executable)) + // Compare venv names against worktree .venv file + let venv_ordering = + wr_venv + .as_ref() + .map_or(Ordering::Equal, |venv| match (&lhs.name, &rhs.name) { + (Some(l), Some(r)) => (r == venv).cmp(&(l == venv)), + (Some(l), None) if l == venv => Ordering::Less, + (None, Some(r)) if r == venv => Ordering::Greater, + _ => Ordering::Equal, + }); + + // Compare project paths against worktree root + let proj_ordering = || match (&lhs.project, &rhs.project) { + (Some(l), Some(r)) => (r == &wr).cmp(&(l == &wr)), + (Some(l), None) if l == &wr => Ordering::Less, + (None, Some(r)) if r == &wr => Ordering::Greater, + _ => Ordering::Equal, + }; + + // Compare environment priorities + let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind)); + + // Compare conda prefixes + let conda_ordering = || { + if lhs.kind == Some(PythonEnvironmentKind::Conda) { + environment + .get_env_var("CONDA_PREFIX".to_string()) + .map(|conda_prefix| { + let is_match = |exe: &Option| { + exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix)) + }; + match (is_match(&lhs.executable), is_match(&rhs.executable)) { + (true, false) => Ordering::Less, + (false, true) => Ordering::Greater, + _ => Ordering::Equal, + } + }) + .unwrap_or(Ordering::Equal) + } else { + Ordering::Equal + } + }; + + // Compare Python executables + let exe_ordering = || lhs.executable.cmp(&rhs.executable); + + venv_ordering + .then_with(proj_ordering) + .then_with(priority_ordering) + .then_with(conda_ordering) + .then_with(exe_ordering) }); let mut toolchains: Vec<_> = toolchains