python: Add recognition of worktree -> venv links (#26759)

Elvis Pranskevichus created

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

Change summary

crates/languages/src/python.rs | 99 +++++++++++++++++++++++++++--------
1 file changed, 76 insertions(+), 23 deletions(-)

Detailed changes

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<PythonEnvironmentKind>) -> usize {
     }
 }
 
+/// Return the name of environment declared in <worktree-root/.venv.
+///
+/// https://virtualfish.readthedocs.io/en/latest/plugins.html#auto-activation-auto-activation
+fn get_worktree_venv_declaration(worktree_root: &Path) -> Option<String> {
+    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 (<workdir>/.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<PathBuf>| {
-                                    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<PathBuf>| {
+                                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