python: Init venv/virtualenv activation scripts during list/resolve (#40816)

This means that existence of activation scripts for venv/virtualenv will
be checked locally either on the host if editing locally, or the remote
by the remote proxy if editing a remote project.

Closes https://github.com/zed-industries/zed/issues/40263

Release Notes:

- N/A

crates/language/src/toolchain.rs

@@ -98,6 +98,7 @@ pub trait ToolchainLister: Send + Sync + 'static {
         worktree_root: PathBuf,
         subroot_relative_path: Arc<RelPath>,
         project_env: Option<HashMap<String, String>>,
+        fs: &dyn Fs,
     ) -> ToolchainList;
 
     /// Given a user-created toolchain, resolve lister-specific details.
@@ -106,14 +107,11 @@ pub trait ToolchainLister: Send + Sync + 'static {
         &self,
         path: PathBuf,
         project_env: Option<HashMap<String, String>>,
+        fs: &dyn Fs,
     ) -> anyhow::Result<Toolchain>;
 
-    async fn activation_script(
-        &self,
-        toolchain: &Toolchain,
-        shell: ShellKind,
-        fs: &dyn Fs,
-    ) -> Vec<String>;
+    fn activation_script(&self, toolchain: &Toolchain, shell: ShellKind) -> Vec<String>;
+
     /// Returns various "static" bits of information about this toolchain lister. This function should be pure.
     fn meta(&self) -> ToolchainMetadata;
 }

crates/languages/src/python.rs

@@ -19,6 +19,7 @@ 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::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use smol::lock::OnceCell;
 use std::cmp::Ordering;
@@ -39,6 +40,14 @@ use std::{
 use task::{ShellKind, TaskTemplate, TaskTemplates, VariableName};
 use util::{ResultExt, maybe};
 
+#[derive(Debug, Serialize, Deserialize)]
+pub(crate) struct PythonToolchainData {
+    #[serde(flatten)]
+    environment: PythonEnvironment,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    activation_scripts: Option<HashMap<ShellKind, PathBuf>>,
+}
+
 pub(crate) struct PyprojectTomlManifestProvider;
 
 impl ManifestProvider for PyprojectTomlManifestProvider {
@@ -165,11 +174,12 @@ impl LspAdapter for TyLspAdapter {
             })?
             .unwrap_or_else(|| json!({}));
         if let Some(toolchain) = toolchain.and_then(|toolchain| {
-            serde_json::from_value::<PythonEnvironment>(toolchain.as_json).ok()
+            serde_json::from_value::<PythonToolchainData>(toolchain.as_json).ok()
         }) {
             _ = maybe!({
-                let uri = url::Url::from_file_path(toolchain.executable?).ok()?;
-                let sys_prefix = toolchain.prefix.clone()?;
+                let uri =
+                    url::Url::from_file_path(toolchain.environment.executable.as_ref()?).ok()?;
+                let sys_prefix = toolchain.environment.prefix.clone()?;
                 let environment = json!({
                     "executable": {
                         "uri": uri,
@@ -474,9 +484,8 @@ impl LspAdapter for PyrightLspAdapter {
 
             // If we have a detected toolchain, configure Pyright to use it
             if let Some(toolchain) = toolchain
-                && let Ok(env) = serde_json::from_value::<
-                    pet_core::python_environment::PythonEnvironment,
-                >(toolchain.as_json.clone())
+                && let Ok(env) =
+                    serde_json::from_value::<PythonToolchainData>(toolchain.as_json.clone())
             {
                 if !user_settings.is_object() {
                     user_settings = Value::Object(serde_json::Map::default());
@@ -484,7 +493,7 @@ impl LspAdapter for PyrightLspAdapter {
                 let object = user_settings.as_object_mut().unwrap();
 
                 let interpreter_path = toolchain.path.to_string();
-                if let Some(venv_dir) = env.prefix {
+                if let Some(venv_dir) = &env.environment.prefix {
                     // Set venvPath and venv at the root level
                     // This matches the format of a pyrightconfig.json file
                     if let Some(parent) = venv_dir.parent() {
@@ -1023,6 +1032,7 @@ impl ToolchainLister for PythonToolchainProvider {
         worktree_root: PathBuf,
         subroot_relative_path: Arc<RelPath>,
         project_env: Option<HashMap<String, String>>,
+        fs: &dyn Fs,
     ) -> ToolchainList {
         let env = project_env.unwrap_or_default();
         let environment = EnvironmentApi::from_env(&env);
@@ -1114,13 +1124,16 @@ impl ToolchainLister for PythonToolchainProvider {
                 .then_with(exe_ordering)
         });
 
-        let mut toolchains: Vec<_> = toolchains
-            .into_iter()
-            .filter_map(venv_to_toolchain)
-            .collect();
-        toolchains.dedup();
+        let mut out_toolchains = Vec::new();
+        for toolchain in toolchains {
+            let Some(toolchain) = venv_to_toolchain(toolchain, fs).await else {
+                continue;
+            };
+            out_toolchains.push(toolchain);
+        }
+        out_toolchains.dedup();
         ToolchainList {
-            toolchains,
+            toolchains: out_toolchains,
             default: None,
             groups: Default::default(),
         }
@@ -1139,6 +1152,7 @@ impl ToolchainLister for PythonToolchainProvider {
         &self,
         path: PathBuf,
         env: Option<HashMap<String, String>>,
+        fs: &dyn Fs,
     ) -> anyhow::Result<Toolchain> {
         let env = env.unwrap_or_default();
         let environment = EnvironmentApi::from_env(&env);
@@ -1150,58 +1164,48 @@ impl ToolchainLister for PythonToolchainProvider {
         let toolchain = pet::resolve::resolve_environment(&path, &locators, &environment)
             .context("Could not find a virtual environment in provided path")?;
         let venv = toolchain.resolved.unwrap_or(toolchain.discovered);
-        venv_to_toolchain(venv).context("Could not convert a venv into a toolchain")
+        venv_to_toolchain(venv, fs)
+            .await
+            .context("Could not convert a venv into a toolchain")
     }
 
-    async fn activation_script(
-        &self,
-        toolchain: &Toolchain,
-        shell: ShellKind,
-        fs: &dyn Fs,
-    ) -> Vec<String> {
-        let Ok(toolchain) = serde_json::from_value::<pet_core::python_environment::PythonEnvironment>(
-            toolchain.as_json.clone(),
-        ) else {
+    fn activation_script(&self, toolchain: &Toolchain, shell: ShellKind) -> Vec<String> {
+        let Ok(toolchain) =
+            serde_json::from_value::<PythonToolchainData>(toolchain.as_json.clone())
+        else {
             return vec![];
         };
+
+        log::debug!("(Python) Composing activation script for toolchain {toolchain:?}");
+
         let mut activation_script = vec![];
 
-        match toolchain.kind {
+        match toolchain.environment.kind {
             Some(PythonEnvironmentKind::Conda) => {
-                if let Some(name) = &toolchain.name {
+                if let Some(name) = &toolchain.environment.name {
                     activation_script.push(format!("conda activate {name}"));
                 } else {
                     activation_script.push("conda activate".to_string());
                 }
             }
             Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
-                if let Some(prefix) = &toolchain.prefix {
-                    let activate_keyword = shell.activate_keyword();
-                    let activate_script_name = match shell {
-                        ShellKind::Posix | ShellKind::Rc => "activate",
-                        ShellKind::Csh => "activate.csh",
-                        ShellKind::Tcsh => "activate.csh",
-                        ShellKind::Fish => "activate.fish",
-                        ShellKind::Nushell => "activate.nu",
-                        ShellKind::PowerShell => "activate.ps1",
-                        ShellKind::Cmd => "activate.bat",
-                        ShellKind::Xonsh => "activate.xsh",
-                    };
-                    let path = prefix.join(BINARY_DIR).join(activate_script_name);
-
-                    if let Some(quoted) = shell.try_quote(&path.to_string_lossy())
-                        && fs.is_file(&path).await
-                    {
-                        activation_script.push(format!("{activate_keyword} {quoted}"));
+                if let Some(activation_scripts) = &toolchain.activation_scripts {
+                    if let Some(activate_script_path) = activation_scripts.get(&shell) {
+                        let activate_keyword = shell.activate_keyword();
+                        if let Some(quoted) =
+                            shell.try_quote(&activate_script_path.to_string_lossy())
+                        {
+                            activation_script.push(format!("{activate_keyword} {quoted}"));
+                        }
                     }
                 }
             }
             Some(PythonEnvironmentKind::Pyenv) => {
-                let Some(manager) = toolchain.manager else {
+                let Some(manager) = &toolchain.environment.manager else {
                     return vec![];
                 };
-                let version = toolchain.version.as_deref().unwrap_or("system");
-                let pyenv = manager.executable;
+                let version = toolchain.environment.version.as_deref().unwrap_or("system");
+                let pyenv = &manager.executable;
                 let pyenv = pyenv.display();
                 activation_script.extend(match shell {
                     ShellKind::Fish => Some(format!("\"{pyenv}\" shell - fish {version}")),
@@ -1221,7 +1225,7 @@ impl ToolchainLister for PythonToolchainProvider {
     }
 }
 
-fn venv_to_toolchain(venv: PythonEnvironment) -> Option<Toolchain> {
+async fn venv_to_toolchain(venv: PythonEnvironment, fs: &dyn Fs) -> Option<Toolchain> {
     let mut name = String::from("Python");
     if let Some(ref version) = venv.version {
         _ = write!(name, " {version}");
@@ -1238,14 +1242,61 @@ fn venv_to_toolchain(venv: PythonEnvironment) -> Option<Toolchain> {
         _ = write!(name, " {nk}");
     }
 
+    let mut activation_scripts = HashMap::default();
+    match venv.kind {
+        Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
+            resolve_venv_activation_scripts(&venv, fs, &mut activation_scripts).await
+        }
+        _ => {}
+    }
+    let data = PythonToolchainData {
+        environment: venv,
+        activation_scripts: Some(activation_scripts),
+    };
+
     Some(Toolchain {
         name: name.into(),
-        path: venv.executable.as_ref()?.to_str()?.to_owned().into(),
+        path: data
+            .environment
+            .executable
+            .as_ref()?
+            .to_str()?
+            .to_owned()
+            .into(),
         language_name: LanguageName::new("Python"),
-        as_json: serde_json::to_value(venv).ok()?,
+        as_json: serde_json::to_value(data).ok()?,
     })
 }
 
+async fn resolve_venv_activation_scripts(
+    venv: &PythonEnvironment,
+    fs: &dyn Fs,
+    activation_scripts: &mut HashMap<ShellKind, PathBuf>,
+) {
+    log::debug!("(Python) Resolving activation scripts for venv toolchain {venv:?}");
+    if let Some(prefix) = &venv.prefix {
+        for (shell_kind, script_name) in &[
+            (ShellKind::Posix, "activate"),
+            (ShellKind::Rc, "activate"),
+            (ShellKind::Csh, "activate.csh"),
+            (ShellKind::Tcsh, "activate.csh"),
+            (ShellKind::Fish, "activate.fish"),
+            (ShellKind::Nushell, "activate.nu"),
+            (ShellKind::PowerShell, "activate.ps1"),
+            (ShellKind::Cmd, "activate.bat"),
+            (ShellKind::Xonsh, "activate.xsh"),
+        ] {
+            let path = prefix.join(BINARY_DIR).join(script_name);
+
+            log::debug!("Trying path: {}", path.display());
+
+            if fs.is_file(&path).await {
+                activation_scripts.insert(*shell_kind, path);
+            }
+        }
+    }
+}
+
 pub struct EnvironmentApi<'a> {
     global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
     project_env: &'a HashMap<String, String>,

crates/project/src/project.rs

@@ -1060,6 +1060,7 @@ impl Project {
                     worktree_store.clone(),
                     environment.clone(),
                     manifest_tree.clone(),
+                    fs.clone(),
                     cx,
                 )
             });

crates/project/src/project_tests.rs

@@ -9653,6 +9653,7 @@ fn python_lang(fs: Arc<FakeFs>) -> Arc<Language> {
             worktree_root: PathBuf,
             subroot_relative_path: Arc<RelPath>,
             _: Option<HashMap<String, String>>,
+            _: &dyn Fs,
         ) -> ToolchainList {
             // This lister will always return a path .venv directories within ancestors
             let ancestors = subroot_relative_path.ancestors().collect::<Vec<_>>();
@@ -9677,6 +9678,7 @@ fn python_lang(fs: Arc<FakeFs>) -> Arc<Language> {
             &self,
             _: PathBuf,
             _: Option<HashMap<String, String>>,
+            _: &dyn Fs,
         ) -> anyhow::Result<Toolchain> {
             Err(anyhow::anyhow!("Not implemented"))
         }
@@ -9689,7 +9691,7 @@ fn python_lang(fs: Arc<FakeFs>) -> Arc<Language> {
                 manifest_name: ManifestName::from(SharedString::new_static("pyproject.toml")),
             }
         }
-        async fn activation_script(&self, _: &Toolchain, _: ShellKind, _: &dyn Fs) -> Vec<String> {
+        fn activation_script(&self, _: &Toolchain, _: ShellKind) -> Vec<String> {
             vec![]
         }
     }

crates/project/src/terminals.rs

@@ -120,7 +120,6 @@ impl Project {
             .map(|p| self.active_toolchain(p, LanguageName::new("Python"), cx))
             .collect::<Vec<_>>();
         let lang_registry = self.languages.clone();
-        let fs = self.fs.clone();
         cx.spawn(async move |project, cx| {
             let shell_kind = ShellKind::new(&shell, is_windows);
             let activation_script = maybe!(async {
@@ -133,11 +132,7 @@ impl Project {
                         .await
                         .ok();
                     let lister = language?.toolchain_lister();
-                    return Some(
-                        lister?
-                            .activation_script(&toolchain, shell_kind, fs.as_ref())
-                            .await,
-                    );
+                    return Some(lister?.activation_script(&toolchain, shell_kind));
                 }
                 None
             })
@@ -347,7 +342,6 @@ impl Project {
         let shell_kind = ShellKind::new(&shell, self.path_style(cx).is_windows());
 
         let lang_registry = self.languages.clone();
-        let fs = self.fs.clone();
         cx.spawn(async move |project, cx| {
             let activation_script = maybe!(async {
                 for toolchain in toolchains {
@@ -359,11 +353,7 @@ impl Project {
                         .await
                         .ok();
                     let lister = language?.toolchain_lister();
-                    return Some(
-                        lister?
-                            .activation_script(&toolchain, shell_kind, fs.as_ref())
-                            .await,
-                    );
+                    return Some(lister?.activation_script(&toolchain, shell_kind));
                 }
                 None
             })

crates/project/src/toolchain_store.rs

@@ -4,6 +4,7 @@ use anyhow::{Context as _, Result, bail};
 
 use async_trait::async_trait;
 use collections::{BTreeMap, IndexSet};
+use fs::Fs;
 use gpui::{
     App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Subscription, Task, WeakEntity,
 };
@@ -60,6 +61,7 @@ impl ToolchainStore {
         worktree_store: Entity<WorktreeStore>,
         project_environment: Entity<ProjectEnvironment>,
         manifest_tree: Entity<ManifestTree>,
+        fs: Arc<dyn Fs>,
         cx: &mut Context<Self>,
     ) -> Self {
         let entity = cx.new(|_| LocalToolchainStore {
@@ -68,6 +70,7 @@ impl ToolchainStore {
             project_environment,
             active_toolchains: Default::default(),
             manifest_tree,
+            fs,
         });
         let _sub = cx.subscribe(&entity, |_, _, e: &ToolchainStoreEvent, cx| {
             cx.emit(e.clone())
@@ -397,6 +400,7 @@ pub struct LocalToolchainStore {
     project_environment: Entity<ProjectEnvironment>,
     active_toolchains: BTreeMap<(WorktreeId, LanguageName), BTreeMap<Arc<RelPath>, Toolchain>>,
     manifest_tree: Entity<ManifestTree>,
+    fs: Arc<dyn Fs>,
 }
 
 #[async_trait(?Send)]
@@ -485,6 +489,7 @@ impl LocalToolchainStore {
         let registry = self.languages.clone();
 
         let manifest_tree = self.manifest_tree.downgrade();
+        let fs = self.fs.clone();
 
         let environment = self.project_environment.clone();
         cx.spawn(async move |this, cx| {
@@ -534,7 +539,12 @@ impl LocalToolchainStore {
             cx.background_spawn(async move {
                 Some((
                     toolchains
-                        .list(worktree_root, relative_path.path.clone(), project_env)
+                        .list(
+                            worktree_root,
+                            relative_path.path.clone(),
+                            project_env,
+                            fs.as_ref(),
+                        )
                         .await,
                     relative_path.path,
                 ))
@@ -568,6 +578,7 @@ impl LocalToolchainStore {
     ) -> Task<Result<Toolchain>> {
         let registry = self.languages.clone();
         let environment = self.project_environment.clone();
+        let fs = self.fs.clone();
         cx.spawn(async move |_, cx| {
             let language = cx
                 .background_spawn(registry.language_for_name(&language_name.0))
@@ -586,8 +597,12 @@ impl LocalToolchainStore {
                     )
                 })?
                 .await;
-            cx.background_spawn(async move { toolchain_lister.resolve(path, project_env).await })
-                .await
+            cx.background_spawn(async move {
+                toolchain_lister
+                    .resolve(path, project_env, fs.as_ref())
+                    .await
+            })
+            .await
         })
     }
 }

crates/remote_server/src/headless_project.rs

@@ -102,6 +102,7 @@ impl HeadlessProject {
                 worktree_store.clone(),
                 environment.clone(),
                 manifest_tree.clone(),
+                fs.clone(),
                 cx,
             )
         });

crates/util/src/shell.rs

@@ -1,6 +1,7 @@
+use serde::{Deserialize, Serialize};
 use std::{borrow::Cow, fmt, path::Path, sync::LazyLock};
 
-#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum ShellKind {
     #[default]
     Posix,