Check if virtual environment is in worktree root (#37510)

George Waters created 1 month ago

The problem from issue #37509 comes from local virtual environments
created with certain approaches (including the 'simple' way of `python
-m venv`) not having a `.project` file with the path to the project's
root directory. When the toolchains are sorted, a virtual environment in
the project is not treated as being for that project and therefore is
not prioritized.

With this change, if a toolchain does not have a `project` associated
with it, we check to see if it is a virtual environment, and if it is we
use its parent directory as the `project`. This will make it the top
priority (i.e. the default) if there are no other virtual environments
for a project, which is what should be expected.

Closes #37509

Release Notes:

- Improved python toolchain prioritization of local virtual
environments.

Change summary

Cargo.lock                     |  1 
Cargo.toml                     |  1 
crates/languages/Cargo.toml    |  1 
crates/languages/src/python.rs | 50 ++++++++++++++++++++++++++++++++---
4 files changed, 48 insertions(+), 5 deletions(-)

Detailed changes

Cargo.lock

@@ -9321,6 +9321,7 @@ dependencies = [
  "pet-fs",
  "pet-poetry",
  "pet-reporter",
+ "pet-virtualenv",
  "pretty_assertions",
  "project",
  "regex",

Cargo.toml

@@ -584,6 +584,7 @@ pet-fs = { git = "https://github.com/microsoft/python-environment-tools.git", re
 pet-pixi = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
 pet-poetry = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
 pet-reporter = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
+pet-virtualenv = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
 portable-pty = "0.9.0"
 postage = { version = "0.5", features = ["futures-traits"] }
 pretty_assertions = { version = "1.3.0", features = ["unstable"] }

crates/languages/Cargo.toml

@@ -57,6 +57,7 @@ pet-core.workspace = true
 pet-fs.workspace = true
 pet-poetry.workspace = true
 pet-reporter.workspace = true
+pet-virtualenv.workspace = true
 pet.workspace = true
 project.workspace = true
 regex.workspace = true

crates/languages/src/python.rs

@@ -16,6 +16,7 @@ use node_runtime::{NodeRuntime, VersionStrategy};
 use pet_core::Configuration;
 use pet_core::os_environment::Environment;
 use pet_core::python_environment::{PythonEnvironment, PythonEnvironmentKind};
+use pet_virtualenv::is_virtualenv_dir;
 use project::Fs;
 use project::lsp_store::language_server_settings;
 use serde_json::{Value, json};
@@ -900,6 +901,21 @@ fn python_module_name_from_relative_path(relative_path: &str) -> String {
         .to_string()
 }
 
+fn is_python_env_global(k: &PythonEnvironmentKind) -> bool {
+    matches!(
+        k,
+        PythonEnvironmentKind::Homebrew
+            | PythonEnvironmentKind::Pyenv
+            | PythonEnvironmentKind::GlobalPaths
+            | PythonEnvironmentKind::MacPythonOrg
+            | PythonEnvironmentKind::MacCommandLineTools
+            | PythonEnvironmentKind::LinuxGlobal
+            | PythonEnvironmentKind::MacXCode
+            | PythonEnvironmentKind::WindowsStore
+            | PythonEnvironmentKind::WindowsRegistry
+    )
+}
+
 fn python_env_kind_display(k: &PythonEnvironmentKind) -> &'static str {
     match k {
         PythonEnvironmentKind::Conda => "Conda",
@@ -966,6 +982,26 @@ async fn get_worktree_venv_declaration(worktree_root: &Path) -> Option<String> {
     Some(venv_name.trim().to_string())
 }
 
+fn get_venv_parent_dir(env: &PythonEnvironment) -> Option<PathBuf> {
+    // If global, we aren't a virtual environment
+    if let Some(kind) = env.kind
+        && is_python_env_global(&kind)
+    {
+        return None;
+    }
+
+    // Check to be sure we are a virtual environment using pet's most generic
+    // virtual environment type, VirtualEnv
+    let venv = env
+        .executable
+        .as_ref()
+        .and_then(|p| p.parent())
+        .and_then(|p| p.parent())
+        .filter(|p| is_virtualenv_dir(p))?;
+
+    venv.parent().map(|parent| parent.to_path_buf())
+}
+
 #[async_trait]
 impl ToolchainLister for PythonToolchainProvider {
     async fn list(
@@ -1025,11 +1061,15 @@ impl ToolchainLister for PythonToolchainProvider {
                     });
 
             // 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,
+            let proj_ordering = || {
+                let lhs_project = lhs.project.clone().or_else(|| get_venv_parent_dir(lhs));
+                let rhs_project = rhs.project.clone().or_else(|| get_venv_parent_dir(rhs));
+                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