repl: List python environments first

crates/repl/src/kernels/mod.rs

@@ -22,11 +22,39 @@ pub trait KernelSession: Sized {
     fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>);
 }
 
+#[derive(Debug, Clone)]
+pub struct PythonEnvKernelSpecification {
+    pub name: String,
+    pub path: PathBuf,
+    pub kernelspec: JupyterKernelspec,
+    pub has_ipykernel: bool,
+    /// Display label for the environment type: "venv", "Conda", "Pyenv", etc.
+    pub environment_kind: Option<String>,
+}
+
+impl PartialEq for PythonEnvKernelSpecification {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name && self.path == other.path
+    }
+}
+
+impl Eq for PythonEnvKernelSpecification {}
+
+impl PythonEnvKernelSpecification {
+    pub fn as_local_spec(&self) -> LocalKernelSpecification {
+        LocalKernelSpecification {
+            name: self.name.clone(),
+            path: self.path.clone(),
+            kernelspec: self.kernelspec.clone(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum KernelSpecification {
     Remote(RemoteKernelSpecification),
     Jupyter(LocalKernelSpecification),
-    PythonEnv(LocalKernelSpecification),
+    PythonEnv(PythonEnvKernelSpecification),
 }
 
 impl KernelSpecification {
@@ -41,7 +69,11 @@ impl KernelSpecification {
     pub fn type_name(&self) -> SharedString {
         match self {
             Self::Jupyter(_) => "Jupyter".into(),
-            Self::PythonEnv(_) => "Python Environment".into(),
+            Self::PythonEnv(spec) => SharedString::from(
+                spec.environment_kind
+                    .clone()
+                    .unwrap_or_else(|| "Python Environment".to_string()),
+            ),
             Self::Remote(_) => "Remote".into(),
         }
     }
@@ -62,6 +94,24 @@ impl KernelSpecification {
         })
     }
 
+    pub fn has_ipykernel(&self) -> bool {
+        match self {
+            Self::Jupyter(_) | Self::Remote(_) => true,
+            Self::PythonEnv(spec) => spec.has_ipykernel,
+        }
+    }
+
+    pub fn environment_kind_label(&self) -> Option<SharedString> {
+        match self {
+            Self::PythonEnv(spec) => spec
+                .environment_kind
+                .as_ref()
+                .map(|kind| SharedString::from(kind.clone())),
+            Self::Jupyter(_) => Some("Jupyter".into()),
+            Self::Remote(_) => Some("Remote".into()),
+        }
+    }
+
     pub fn icon(&self, cx: &App) -> Icon {
         let lang_name = match self {
             Self::Jupyter(spec) => spec.kernelspec.language.clone(),
@@ -76,6 +126,33 @@ impl KernelSpecification {
     }
 }
 
+fn extract_environment_kind(toolchain_json: &serde_json::Value) -> Option<String> {
+    let kind_str = toolchain_json.get("kind")?.as_str()?;
+    let label = match kind_str {
+        "Conda" => "Conda",
+        "Pixi" => "pixi",
+        "Homebrew" => "Homebrew",
+        "Pyenv" => "global (Pyenv)",
+        "GlobalPaths" => "global",
+        "PyenvVirtualEnv" => "Pyenv",
+        "Pipenv" => "Pipenv",
+        "Poetry" => "Poetry",
+        "MacPythonOrg" => "global (Python.org)",
+        "MacCommandLineTools" => "global (Command Line Tools for Xcode)",
+        "LinuxGlobal" => "global",
+        "MacXCode" => "global (Xcode)",
+        "Venv" => "venv",
+        "VirtualEnv" => "virtualenv",
+        "VirtualEnvWrapper" => "virtualenvwrapper",
+        "WindowsStore" => "global (Windows Store)",
+        "WindowsRegistry" => "global (Windows Registry)",
+        "Uv" => "uv",
+        "UvWorkspace" => "uv (Workspace)",
+        _ => kind_str,
+    };
+    Some(label.to_string())
+}
+
 pub fn python_env_kernel_specifications(
     project: &Entity<Project>,
     worktree_id: WorktreeId,
@@ -111,46 +188,41 @@ pub fn python_env_kernel_specifications(
             .map(|toolchain| {
                 background_executor.spawn(async move {
                     let python_path = toolchain.path.to_string();
+                    let environment_kind = extract_environment_kind(&toolchain.as_json);
 
-                    // Check if ipykernel is installed
-                    let ipykernel_check = util::command::new_smol_command(&python_path)
+                    let has_ipykernel = util::command::new_smol_command(&python_path)
                         .args(&["-c", "import ipykernel"])
                         .output()
-                        .await;
-
-                    if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
-                        // Create a default kernelspec for this environment
-                        let default_kernelspec = JupyterKernelspec {
-                            argv: vec![
-                                python_path.clone(),
-                                "-m".to_string(),
-                                "ipykernel_launcher".to_string(),
-                                "-f".to_string(),
-                                "{connection_file}".to_string(),
-                            ],
-                            display_name: toolchain.name.to_string(),
-                            language: "python".to_string(),
-                            interrupt_mode: None,
-                            metadata: None,
-                            env: None,
-                        };
-
-                        Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
-                            name: toolchain.name.to_string(),
-                            path: PathBuf::from(&python_path),
-                            kernelspec: default_kernelspec,
-                        }))
-                    } else {
-                        None
-                    }
+                        .await
+                        .map(|output| output.status.success())
+                        .unwrap_or(false);
+
+                    let kernelspec = JupyterKernelspec {
+                        argv: vec![
+                            python_path.clone(),
+                            "-m".to_string(),
+                            "ipykernel_launcher".to_string(),
+                            "-f".to_string(),
+                            "{connection_file}".to_string(),
+                        ],
+                        display_name: toolchain.name.to_string(),
+                        language: "python".to_string(),
+                        interrupt_mode: None,
+                        metadata: None,
+                        env: None,
+                    };
+
+                    KernelSpecification::PythonEnv(PythonEnvKernelSpecification {
+                        name: toolchain.name.to_string(),
+                        path: PathBuf::from(&python_path),
+                        kernelspec,
+                        has_ipykernel,
+                        environment_kind,
+                    })
                 })
             });
 
-        let kernel_specs = futures::future::join_all(kernelspecs)
-            .await
-            .into_iter()
-            .flatten()
-            .collect();
+        let kernel_specs = futures::future::join_all(kernelspecs).await;
 
         anyhow::Ok(kernel_specs)
     }

crates/repl/src/notebook/notebook_ui.rs

@@ -414,8 +414,7 @@ impl NotebookEditor {
         });
 
         let kernel_task = match spec {
-            KernelSpecification::Jupyter(local_spec)
-            | KernelSpecification::PythonEnv(local_spec) => NativeRunningKernel::new(
+            KernelSpecification::Jupyter(local_spec) => NativeRunningKernel::new(
                 local_spec,
                 entity_id,
                 working_directory,
@@ -424,6 +423,15 @@ impl NotebookEditor {
                 window,
                 cx,
             ),
+            KernelSpecification::PythonEnv(env_spec) => NativeRunningKernel::new(
+                env_spec.as_local_spec(),
+                entity_id,
+                working_directory,
+                fs,
+                view,
+                window,
+                cx,
+            ),
             KernelSpecification::Remote(remote_spec) => {
                 RemoteRunningKernel::new(remote_spec, working_directory, view, window, cx)
             }

crates/repl/src/repl.rs

@@ -17,13 +17,13 @@ use project::Fs;
 pub use runtimelib::ExecutionState;
 
 pub use crate::jupyter_settings::JupyterSettings;
-pub use crate::kernels::{Kernel, KernelSpecification, KernelStatus};
+pub use crate::kernels::{Kernel, KernelSpecification, KernelStatus, PythonEnvKernelSpecification};
 pub use crate::repl_editor::*;
 pub use crate::repl_sessions_ui::{
     ClearOutputs, Interrupt, ReplSessionsPage, Restart, Run, Sessions, Shutdown,
 };
 pub use crate::repl_settings::ReplSettings;
-use crate::repl_store::ReplStore;
+pub use crate::repl_store::ReplStore;
 pub use crate::session::Session;
 
 pub const KERNEL_DOCS_URL: &str = "https://zed.dev/docs/repl#changing-kernels";

crates/repl/src/repl_editor.rs

@@ -8,7 +8,12 @@ use editor::{Editor, MultiBufferOffset};
 use gpui::{App, Entity, WeakEntity, Window, prelude::*};
 use language::{BufferSnapshot, Language, LanguageName, Point};
 use project::{ProjectItem as _, WorktreeId};
+use workspace::{
+    Workspace,
+    notifications::{NotificationId, NotificationSource},
+};
 
+use crate::kernels::PythonEnvKernelSpecification;
 use crate::repl_store::ReplStore;
 use crate::session::SessionEvent;
 use crate::{
@@ -72,6 +77,115 @@ pub fn assign_kernelspec(
     Ok(())
 }
 
+pub fn install_ipykernel_and_assign(
+    kernel_specification: KernelSpecification,
+    weak_editor: WeakEntity<Editor>,
+    window: &mut Window,
+    cx: &mut App,
+) -> Result<()> {
+    let KernelSpecification::PythonEnv(ref env_spec) = kernel_specification else {
+        return assign_kernelspec(kernel_specification, weak_editor, window, cx);
+    };
+
+    let python_path = env_spec.path.clone();
+    let env_name = env_spec.name.clone();
+    let env_spec = env_spec.clone();
+
+    struct IpykernelInstall;
+    let notification_id = NotificationId::unique::<IpykernelInstall>();
+
+    let workspace = Workspace::for_window(window, cx);
+    if let Some(workspace) = &workspace {
+        workspace.update(cx, |workspace, cx| {
+            workspace.show_toast(
+                workspace::Toast::new(
+                    notification_id.clone(),
+                    format!("Installing ipykernel in {}...", env_name),
+                ),
+                NotificationSource::Project,
+                cx,
+            );
+        });
+    }
+
+    let weak_workspace = workspace.map(|w| w.downgrade());
+    let window_handle = window.window_handle();
+
+    let install_task = cx.background_spawn(async move {
+        let output = util::command::new_smol_command(python_path.to_string_lossy().as_ref())
+            .args(&["-m", "pip", "install", "ipykernel"])
+            .output()
+            .await
+            .context("failed to run pip install ipykernel")?;
+
+        if output.status.success() {
+            anyhow::Ok(())
+        } else {
+            let stderr = String::from_utf8_lossy(&output.stderr);
+            anyhow::bail!("{}", stderr.lines().last().unwrap_or("unknown error"))
+        }
+    });
+
+    cx.spawn(async move |cx| {
+        let result = install_task.await;
+
+        match result {
+            Ok(()) => {
+                if let Some(weak_workspace) = &weak_workspace {
+                    weak_workspace
+                        .update(cx, |workspace, cx| {
+                            workspace.dismiss_toast(&notification_id, cx);
+                            workspace.show_toast(
+                                workspace::Toast::new(
+                                    notification_id.clone(),
+                                    format!("ipykernel installed in {}", env_name),
+                                )
+                                .autohide(),
+                                NotificationSource::Project,
+                                cx,
+                            );
+                        })
+                        .ok();
+                }
+
+                window_handle
+                    .update(cx, |_, window, cx| {
+                        let updated_spec =
+                            KernelSpecification::PythonEnv(PythonEnvKernelSpecification {
+                                has_ipykernel: true,
+                                ..env_spec
+                            });
+                        assign_kernelspec(updated_spec, weak_editor, window, cx).ok();
+                    })
+                    .ok();
+            }
+            Err(error) => {
+                if let Some(weak_workspace) = &weak_workspace {
+                    weak_workspace
+                        .update(cx, |workspace, cx| {
+                            workspace.dismiss_toast(&notification_id, cx);
+                            workspace.show_toast(
+                                workspace::Toast::new(
+                                    notification_id.clone(),
+                                    format!(
+                                        "Failed to install ipykernel in {}: {}",
+                                        env_name, error
+                                    ),
+                                ),
+                                NotificationSource::Project,
+                                cx,
+                            );
+                        })
+                        .ok();
+                }
+            }
+        }
+    })
+    .detach();
+
+    Ok(())
+}
+
 pub fn run(
     editor: WeakEntity<Editor>,
     move_down: bool,