python: Improve sorting order of toolchains to give higher precedence to project-local virtual environments that are within current subproject (#44141)

Closes #44090

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>

Release Notes:

- python: Improved sorting order of toolchains in monorepos with
multiple local virtual environments.
- python: Fixed toolchain selector not having an active toolchain
selected on open.

---------

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>
Co-authored-by: Smit <smit@zed.dev>

crates/languages/src/python.rs

@@ -23,7 +23,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use settings::Settings;
 use smol::lock::OnceCell;
-use std::cmp::Ordering;
+use std::cmp::{Ordering, Reverse};
 use std::env::consts;
 use terminal::terminal_settings::TerminalSettings;
 use util::command::new_smol_command;
@@ -1101,13 +1101,33 @@ fn get_venv_parent_dir(env: &PythonEnvironment) -> Option<PathBuf> {
     venv.parent().map(|parent| parent.to_path_buf())
 }
 
-fn wr_distance(wr: &PathBuf, venv: Option<&PathBuf>) -> usize {
+// How far is this venv from the root of our current project?
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
+enum SubprojectDistance {
+    WithinSubproject(Reverse<usize>),
+    WithinWorktree(Reverse<usize>),
+    NotInWorktree,
+}
+
+fn wr_distance(
+    wr: &PathBuf,
+    subroot_relative_path: &RelPath,
+    venv: Option<&PathBuf>,
+) -> SubprojectDistance {
     if let Some(venv) = venv
         && let Ok(p) = venv.strip_prefix(wr)
     {
-        p.components().count()
+        if subroot_relative_path.components().next().is_some()
+            && let Ok(distance) = p
+                .strip_prefix(subroot_relative_path.as_std_path())
+                .map(|p| p.components().count())
+        {
+            SubprojectDistance::WithinSubproject(Reverse(distance))
+        } else {
+            SubprojectDistance::WithinWorktree(Reverse(p.components().count()))
+        }
     } else {
-        usize::MAX
+        SubprojectDistance::NotInWorktree
     }
 }
 
@@ -1170,11 +1190,14 @@ impl ToolchainLister for PythonToolchainProvider {
                     });
 
             // Compare project paths against worktree root
-            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));
-                wr_distance(&wr, lhs_project.as_ref()).cmp(&wr_distance(&wr, rhs_project.as_ref()))
-            };
+            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));
+                    wr_distance(&wr, &subroot_relative_path, lhs_project.as_ref()).cmp(
+                        &wr_distance(&wr, &subroot_relative_path, rhs_project.as_ref()),
+                    )
+                };
 
             // Compare environment priorities
             let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind));

crates/toolchain_selector/src/toolchain_selector.rs

@@ -588,19 +588,20 @@ impl ToolchainSelector {
             .worktree_for_id(worktree_id, cx)?
             .read(cx)
             .abs_path();
-        let workspace_id = workspace.database_id()?;
         let weak = workspace.weak_handle();
         cx.spawn_in(window, async move |workspace, cx| {
-            let active_toolchain = workspace::WORKSPACE_DB
-                .toolchain(
-                    workspace_id,
-                    worktree_id,
-                    relative_path.clone(),
-                    language_name.clone(),
-                )
-                .await
-                .ok()
-                .flatten();
+            let active_toolchain = project
+                .read_with(cx, |this, cx| {
+                    this.active_toolchain(
+                        ProjectPath {
+                            worktree_id,
+                            path: relative_path.clone(),
+                        },
+                        language_name.clone(),
+                        cx,
+                    )
+                })?
+                .await;
             workspace
                 .update_in(cx, |this, window, cx| {
                     this.toggle_modal(window, cx, move |window, cx| {
@@ -618,6 +619,7 @@ impl ToolchainSelector {
                     });
                 })
                 .ok();
+            anyhow::Ok(())
         })
         .detach();
 

crates/workspace/src/persistence.rs

@@ -1656,49 +1656,6 @@ impl WorkspaceDb {
         }
     }
 
-    pub async fn toolchain(
-        &self,
-        workspace_id: WorkspaceId,
-        worktree_id: WorktreeId,
-        relative_worktree_path: Arc<RelPath>,
-        language_name: LanguageName,
-    ) -> Result<Option<Toolchain>> {
-        self.write(move |this| {
-            let mut select = this
-                .select_bound(sql!(
-                    SELECT
-                        name, path, raw_json
-                    FROM toolchains
-                    WHERE
-                        workspace_id = ? AND
-                        language_name = ? AND
-                        worktree_id = ? AND
-                        relative_worktree_path = ?
-                ))
-                .context("select toolchain")?;
-
-            let toolchain: Vec<(String, String, String)> = select((
-                workspace_id,
-                language_name.as_ref().to_string(),
-                worktree_id.to_usize(),
-                relative_worktree_path.as_unix_str().to_string(),
-            ))?;
-
-            Ok(toolchain
-                .into_iter()
-                .next()
-                .and_then(|(name, path, raw_json)| {
-                    Some(Toolchain {
-                        name: name.into(),
-                        path: path.into(),
-                        language_name,
-                        as_json: serde_json::Value::from_str(&raw_json).ok()?,
-                    })
-                }))
-        })
-        .await
-    }
-
     pub(crate) async fn toolchains(
         &self,
         workspace_id: WorkspaceId,