From 317a578f6a4a7cfa34a2350f72e2b328b8bc86bb Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Wed, 11 Feb 2026 23:20:20 -0800 Subject: [PATCH] repl: List python environments first (#48763) This PR completely subsumes https://github.com/zed-industries/zed/pull/46720 image Release Notes: - Added Python Environments to REPL kernel selection - List active toolchain/python environment as the recommended kernel for REPL usage --------- Co-authored-by: Claude Opus 4.5 --- crates/repl/src/components/kernel_options.rs | 425 +++++++++++++----- crates/repl/src/kernels/mod.rs | 144 ++++-- crates/repl/src/notebook/notebook_ui.rs | 12 +- crates/repl/src/repl.rs | 4 +- crates/repl/src/repl_editor.rs | 114 +++++ crates/repl/src/repl_store.rs | 102 ++++- crates/repl/src/session.rs | 12 +- .../zed/src/zed/quick_action_bar/repl_menu.rs | 27 +- 8 files changed, 651 insertions(+), 189 deletions(-) diff --git a/crates/repl/src/components/kernel_options.rs b/crates/repl/src/components/kernel_options.rs index bceefd08cc8f897cc33a3bdb4b1be5c2dc90df10..79a236f7c6478c1bc9b9e48ac17596341dfa8aaf 100644 --- a/crates/repl/src/components/kernel_options.rs +++ b/crates/repl/src/components/kernel_options.rs @@ -2,23 +2,107 @@ use crate::KERNEL_DOCS_URL; use crate::kernels::KernelSpecification; use crate::repl_store::ReplStore; -use gpui::AnyView; -use gpui::DismissEvent; - -use gpui::FontWeight; -use picker::Picker; -use picker::PickerDelegate; +use gpui::{AnyView, DismissEvent, FontWeight, SharedString, Task}; +use picker::{Picker, PickerDelegate}; use project::WorktreeId; - use std::sync::Arc; -use ui::ListItemSpacing; - -use gpui::SharedString; -use gpui::Task; -use ui::{ListItem, PopoverMenu, PopoverMenuHandle, PopoverTrigger, prelude::*}; +use ui::{ListItem, ListItemSpacing, PopoverMenu, PopoverMenuHandle, PopoverTrigger, prelude::*}; type OnSelect = Box; +#[derive(Clone)] +pub enum KernelPickerEntry { + SectionHeader(SharedString), + Kernel { + spec: KernelSpecification, + is_recommended: bool, + }, +} + +fn build_grouped_entries(store: &ReplStore, worktree_id: WorktreeId) -> Vec { + let mut entries = Vec::new(); + let mut recommended_entry: Option = None; + + let mut python_envs = Vec::new(); + let mut jupyter_kernels = Vec::new(); + let mut remote_kernels = Vec::new(); + + for spec in store.kernel_specifications_for_worktree(worktree_id) { + let is_recommended = store.is_recommended_kernel(worktree_id, spec); + + if is_recommended { + recommended_entry = Some(KernelPickerEntry::Kernel { + spec: spec.clone(), + is_recommended: true, + }); + } + + match spec { + KernelSpecification::PythonEnv(_) => { + python_envs.push(KernelPickerEntry::Kernel { + spec: spec.clone(), + is_recommended, + }); + } + KernelSpecification::Jupyter(_) => { + jupyter_kernels.push(KernelPickerEntry::Kernel { + spec: spec.clone(), + is_recommended, + }); + } + KernelSpecification::Remote(_) => { + remote_kernels.push(KernelPickerEntry::Kernel { + spec: spec.clone(), + is_recommended, + }); + } + } + } + + // Sort Python envs: has_ipykernel first, then by name + python_envs.sort_by(|a, b| { + let (spec_a, spec_b) = match (a, b) { + ( + KernelPickerEntry::Kernel { spec: sa, .. }, + KernelPickerEntry::Kernel { spec: sb, .. }, + ) => (sa, sb), + _ => return std::cmp::Ordering::Equal, + }; + spec_b + .has_ipykernel() + .cmp(&spec_a.has_ipykernel()) + .then_with(|| spec_a.name().cmp(&spec_b.name())) + }); + + // Recommended section + if let Some(rec) = recommended_entry { + entries.push(KernelPickerEntry::SectionHeader("Recommended".into())); + entries.push(rec); + } + + // Python Environments section + if !python_envs.is_empty() { + entries.push(KernelPickerEntry::SectionHeader( + "Python Environments".into(), + )); + entries.extend(python_envs); + } + + // Jupyter Kernels section + if !jupyter_kernels.is_empty() { + entries.push(KernelPickerEntry::SectionHeader("Jupyter Kernels".into())); + entries.extend(jupyter_kernels); + } + + // Remote section + if !remote_kernels.is_empty() { + entries.push(KernelPickerEntry::SectionHeader("Remote Servers".into())); + entries.extend(remote_kernels); + } + + entries +} + #[derive(IntoElement)] pub struct KernelSelector where @@ -34,22 +118,13 @@ where } pub struct KernelPickerDelegate { - all_kernels: Vec, - filtered_kernels: Vec, + all_entries: Vec, + filtered_entries: Vec, selected_kernelspec: Option, + selected_index: usize, on_select: OnSelect, } -// Helper function to truncate long paths -fn truncate_path(path: &SharedString, max_length: usize) -> SharedString { - if path.len() <= max_length { - path.to_string().into() - } else { - let truncated = path.chars().rev().take(max_length - 3).collect::(); - format!("...{}", truncated.chars().rev().collect::()).into() - } -} - impl KernelSelector where T: PopoverTrigger + ButtonCommon, @@ -77,26 +152,66 @@ where } } +impl KernelPickerDelegate { + fn first_selectable_index(entries: &[KernelPickerEntry]) -> usize { + entries + .iter() + .position(|e| matches!(e, KernelPickerEntry::Kernel { .. })) + .unwrap_or(0) + } + + fn next_selectable_index(&self, from: usize, direction: i32) -> usize { + let len = self.filtered_entries.len(); + if len == 0 { + return 0; + } + + let mut index = from as i32 + direction; + while index >= 0 && (index as usize) < len { + if matches!( + self.filtered_entries.get(index as usize), + Some(KernelPickerEntry::Kernel { .. }) + ) { + return index as usize; + } + index += direction; + } + + from + } +} + impl PickerDelegate for KernelPickerDelegate { type ListItem = ListItem; fn match_count(&self) -> usize { - self.filtered_kernels.len() + self.filtered_entries.len() } fn selected_index(&self) -> usize { - if let Some(kernelspec) = self.selected_kernelspec.as_ref() { - self.filtered_kernels - .iter() - .position(|k| k == kernelspec) - .unwrap_or(0) - } else { - 0 - } + self.selected_index } fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context>) { - self.selected_kernelspec = self.filtered_kernels.get(ix).cloned(); + if matches!( + self.filtered_entries.get(ix), + Some(KernelPickerEntry::SectionHeader(_)) + ) { + let forward = self.next_selectable_index(ix, 1); + if forward != ix { + self.selected_index = forward; + } else { + self.selected_index = self.next_selectable_index(ix, -1); + } + } else { + self.selected_index = ix; + } + + if let Some(KernelPickerEntry::Kernel { spec, .. }) = + self.filtered_entries.get(self.selected_index) + { + self.selected_kernelspec = Some(spec.clone()); + } cx.notify(); } @@ -110,28 +225,57 @@ impl PickerDelegate for KernelPickerDelegate { _window: &mut Window, _cx: &mut Context>, ) -> Task<()> { - let all_kernels = self.all_kernels.clone(); - if query.is_empty() { - self.filtered_kernels = all_kernels; - return Task::ready(()); + self.filtered_entries = self.all_entries.clone(); + } else { + let query_lower = query.to_lowercase(); + let mut filtered = Vec::new(); + let mut pending_header: Option = None; + + for entry in &self.all_entries { + match entry { + KernelPickerEntry::SectionHeader(_) => { + pending_header = Some(entry.clone()); + } + KernelPickerEntry::Kernel { spec, .. } => { + if spec.name().to_lowercase().contains(&query_lower) { + if let Some(header) = pending_header.take() { + filtered.push(header); + } + filtered.push(entry.clone()); + } + } + } + } + + self.filtered_entries = filtered; } - self.filtered_kernels = if query.is_empty() { - all_kernels - } else { - all_kernels - .into_iter() - .filter(|kernel| kernel.name().to_lowercase().contains(&query.to_lowercase())) - .collect() - }; + self.selected_index = Self::first_selectable_index(&self.filtered_entries); + if let Some(KernelPickerEntry::Kernel { spec, .. }) = + self.filtered_entries.get(self.selected_index) + { + self.selected_kernelspec = Some(spec.clone()); + } Task::ready(()) } + fn separators_after_indices(&self) -> Vec { + let mut separators = Vec::new(); + for (index, entry) in self.filtered_entries.iter().enumerate() { + if matches!(entry, KernelPickerEntry::SectionHeader(_)) && index > 0 { + separators.push(index - 1); + } + } + separators + } + fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context>) { - if let Some(kernelspec) = &self.selected_kernelspec { - (self.on_select)(kernelspec.clone(), window, cx); + if let Some(KernelPickerEntry::Kernel { spec, .. }) = + self.filtered_entries.get(self.selected_index) + { + (self.on_select)(spec.clone(), window, cx); cx.emit(DismissEvent); } } @@ -145,80 +289,107 @@ impl PickerDelegate for KernelPickerDelegate { _: &mut Window, cx: &mut Context>, ) -> Option { - let kernelspec = self.filtered_kernels.get(ix)?; - let is_selected = self.selected_kernelspec.as_ref() == Some(kernelspec); - let icon = kernelspec.icon(cx); - - let (name, kernel_type, path_or_url) = match kernelspec { - KernelSpecification::Jupyter(_) => (kernelspec.name(), "Jupyter", None), - KernelSpecification::PythonEnv(_) => ( - kernelspec.name(), - "Python Env", - Some(truncate_path(&kernelspec.path(), 42)), - ), - KernelSpecification::Remote(_) => ( - kernelspec.name(), - "Remote", - Some(truncate_path(&kernelspec.path(), 42)), + let entry = self.filtered_entries.get(ix)?; + + match entry { + KernelPickerEntry::SectionHeader(title) => Some( + ListItem::new(ix) + .inset(true) + .spacing(ListItemSpacing::Dense) + .selectable(false) + .child( + Label::new(title.clone()) + .size(LabelSize::Small) + .weight(FontWeight::SEMIBOLD) + .color(Color::Muted), + ), ), - }; - - Some( - ListItem::new(ix) - .inset(true) - .spacing(ListItemSpacing::Sparse) - .toggle_state(selected) - .child( - h_flex() - .w_full() - .gap_3() - .child(icon.color(Color::Default).size(IconSize::Medium)) + KernelPickerEntry::Kernel { + spec, + is_recommended, + } => { + let is_currently_selected = self.selected_kernelspec.as_ref() == Some(spec); + let icon = spec.icon(cx); + let has_ipykernel = spec.has_ipykernel(); + + let subtitle = match spec { + KernelSpecification::Jupyter(_) => None, + KernelSpecification::PythonEnv(_) | KernelSpecification::Remote(_) => { + let env_kind = spec.environment_kind_label(); + let path = spec.path(); + match env_kind { + Some(kind) => Some(format!("{} \u{2013} {}", kind, path)), + None => Some(path.to_string()), + } + } + }; + + Some( + ListItem::new(ix) + .inset(true) + .spacing(ListItemSpacing::Sparse) + .toggle_state(selected) .child( - v_flex() - .flex_grow() - .gap_0p5() + h_flex() + .w_full() + .gap_3() + .when(!has_ipykernel, |flex| flex.opacity(0.5)) + .child(icon.color(Color::Default).size(IconSize::Medium)) .child( - h_flex() - .justify_between() + v_flex() + .flex_grow() + .overflow_x_hidden() + .gap_0p5() .child( - div().w_48().text_ellipsis().child( - Label::new(name) - .weight(FontWeight::MEDIUM) - .size(LabelSize::Default), - ), + h_flex() + .gap_1() + .child( + div() + .overflow_x_hidden() + .flex_shrink() + .text_ellipsis() + .child( + Label::new(spec.name()) + .weight(FontWeight::MEDIUM) + .size(LabelSize::Default), + ), + ) + .when(*is_recommended, |flex| { + flex.child( + Label::new("Recommended") + .size(LabelSize::XSmall) + .color(Color::Accent), + ) + }) + .when(!has_ipykernel, |flex| { + flex.child( + Label::new("ipykernel not installed") + .size(LabelSize::XSmall) + .color(Color::Warning), + ) + }), ) - .when_some(path_or_url, |flex, path| { - flex.text_ellipsis().child( - Label::new(path) - .size(LabelSize::Small) - .color(Color::Muted), + .when_some(subtitle, |flex, subtitle| { + flex.child( + div().overflow_x_hidden().text_ellipsis().child( + Label::new(subtitle) + .size(LabelSize::Small) + .color(Color::Muted), + ), ) }), - ) - .child( - h_flex() - .gap_1() - .child( - Label::new(kernelspec.language()) - .size(LabelSize::Small) - .color(Color::Muted), - ) - .child( - Label::new(kernel_type) - .size(LabelSize::Small) - .color(Color::Muted), - ), ), - ), + ) + .when(is_currently_selected, |item| { + item.end_slot( + Icon::new(IconName::Check) + .color(Color::Accent) + .size(IconSize::Small), + ) + }), ) - .when(is_selected, |item| { - item.end_slot( - Icon::new(IconName::Check) - .color(Color::Accent) - .size(IconSize::Small), - ) - }), - ) + } + } } fn render_footer( @@ -254,24 +425,32 @@ where fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement { let store = ReplStore::global(cx).read(cx); - let all_kernels: Vec = store - .kernel_specifications_for_worktree(self.worktree_id) - .cloned() - .collect(); - + let all_entries = build_grouped_entries(store, self.worktree_id); let selected_kernelspec = store.active_kernelspec(self.worktree_id, None, cx); + let selected_index = all_entries + .iter() + .position(|entry| { + if let KernelPickerEntry::Kernel { spec, .. } = entry { + selected_kernelspec.as_ref() == Some(spec) + } else { + false + } + }) + .unwrap_or_else(|| KernelPickerDelegate::first_selectable_index(&all_entries)); let delegate = KernelPickerDelegate { on_select: self.on_select, - all_kernels: all_kernels.clone(), - filtered_kernels: all_kernels, + all_entries: all_entries.clone(), + filtered_entries: all_entries, selected_kernelspec, + selected_index, }; let picker_view = cx.new(|cx| { - Picker::uniform_list(delegate, window, cx) - .width(rems(30.)) - .max_height(Some(rems(20.).into())) + Picker::list(delegate, window, cx) + .list_measure_all() + .width(rems(34.)) + .max_height(Some(rems(24.).into())) }); PopoverMenu::new("kernel-switcher") diff --git a/crates/repl/src/kernels/mod.rs b/crates/repl/src/kernels/mod.rs index aaaaa40765e1fa4c8abc3cab4394f4f91d77a6b6..ceef195f737465afd064790b675e4051786b5aa6 100644 --- a/crates/repl/src/kernels/mod.rs +++ b/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); } +#[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, +} + +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 { + 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 { + 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, 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) } diff --git a/crates/repl/src/notebook/notebook_ui.rs b/crates/repl/src/notebook/notebook_ui.rs index 6c3848046b5c330ebfacccca090d23db50242af6..9fe9811c5ee7e36b4ecd09b79e8d57fb5e995f3f 100644 --- a/crates/repl/src/notebook/notebook_ui.rs +++ b/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) } diff --git a/crates/repl/src/repl.rs b/crates/repl/src/repl.rs index 6c0993169144d7d9c1c9f97a6fb4572173021ccf..be64973e52d4750e4dbb17f944507f245711774b 100644 --- a/crates/repl/src/repl.rs +++ b/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"; diff --git a/crates/repl/src/repl_editor.rs b/crates/repl/src/repl_editor.rs index 85bd8bfb4a070f33f5a75d97d449ab6fcf9d5211..c9cb1d26ce814a947a2fa8bef0503a602747b5c3 100644 --- a/crates/repl/src/repl_editor.rs +++ b/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, + 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::(); + + 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(¬ification_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(¬ification_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, move_down: bool, diff --git a/crates/repl/src/repl_store.rs b/crates/repl/src/repl_store.rs index abff0bdc3aa20beacbd566c512bbbab1064a1610..1fd720d977b91d52a4cc8dc74ee30cdb8e5a2de1 100644 --- a/crates/repl/src/repl_store.rs +++ b/crates/repl/src/repl_store.rs @@ -4,11 +4,12 @@ use std::sync::Arc; use anyhow::{Context as _, Result}; use collections::HashMap; use command_palette_hooks::CommandPaletteFilter; -use gpui::{App, Context, Entity, EntityId, Global, Subscription, Task, prelude::*}; +use gpui::{App, Context, Entity, EntityId, Global, SharedString, Subscription, Task, prelude::*}; use jupyter_websocket_client::RemoteServer; -use language::Language; -use project::{Fs, Project, WorktreeId}; +use language::{Language, LanguageName}; +use project::{Fs, Project, ProjectPath, WorktreeId}; use settings::{Settings, SettingsStore}; +use util::rel_path::RelPath; use crate::kernels::{ Kernel, list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications, @@ -26,6 +27,7 @@ pub struct ReplStore { kernel_specifications: Vec, selected_kernel_for_worktree: HashMap, kernel_specifications_for_worktree: HashMap>, + active_python_toolchain_for_worktree: HashMap, _subscriptions: Vec, } @@ -63,6 +65,7 @@ impl ReplStore { _subscriptions: subscriptions, kernel_specifications_for_worktree: HashMap::default(), selected_kernel_for_worktree: HashMap::default(), + active_python_toolchain_for_worktree: HashMap::default(), }; this.on_enabled_changed(cx); this @@ -76,6 +79,11 @@ impl ReplStore { self.enabled } + pub fn has_python_kernelspecs(&self, worktree_id: WorktreeId) -> bool { + self.kernel_specifications_for_worktree + .contains_key(&worktree_id) + } + pub fn kernel_specifications_for_worktree( &self, worktree_id: WorktreeId, @@ -127,14 +135,29 @@ impl ReplStore { cx: &mut Context, ) -> Task> { let kernel_specifications = python_env_kernel_specifications(project, worktree_id, cx); + let active_toolchain = project.read(cx).active_toolchain( + ProjectPath { + worktree_id, + path: RelPath::empty().into(), + }, + LanguageName::new_static("Python"), + cx, + ); + cx.spawn(async move |this, cx| { let kernel_specifications = kernel_specifications .await .context("getting python kernelspecs")?; + let active_toolchain_path = active_toolchain.await.map(|toolchain| toolchain.path); + this.update(cx, |this, cx| { this.kernel_specifications_for_worktree .insert(worktree_id, kernel_specifications); + if let Some(path) = active_toolchain_path { + this.active_python_toolchain_for_worktree + .insert(worktree_id, path); + } cx.notify(); }) }) @@ -210,21 +233,65 @@ impl ReplStore { .insert(worktree_id, kernelspec); } + pub fn active_python_toolchain_path(&self, worktree_id: WorktreeId) -> Option<&SharedString> { + self.active_python_toolchain_for_worktree.get(&worktree_id) + } + + pub fn is_recommended_kernel( + &self, + worktree_id: WorktreeId, + spec: &KernelSpecification, + ) -> bool { + if let Some(active_path) = self.active_python_toolchain_path(worktree_id) { + spec.path().as_ref() == active_path.as_ref() + } else { + false + } + } + pub fn active_kernelspec( &self, worktree_id: WorktreeId, language_at_cursor: Option>, cx: &App, ) -> Option { - let selected_kernelspec = self.selected_kernel_for_worktree.get(&worktree_id).cloned(); + if let Some(selected) = self.selected_kernel_for_worktree.get(&worktree_id).cloned() { + return Some(selected); + } - if let Some(language_at_cursor) = language_at_cursor { - selected_kernelspec.or_else(|| { - self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx) + let language_at_cursor = language_at_cursor?; + let language_name = language_at_cursor.code_fence_block_name().to_lowercase(); + + // Prefer the recommended (active toolchain) kernel if it has ipykernel + if let Some(active_path) = self.active_python_toolchain_path(worktree_id) { + let recommended = self + .kernel_specifications_for_worktree(worktree_id) + .find(|spec| { + spec.has_ipykernel() + && spec.language().as_ref().to_lowercase() == language_name + && spec.path().as_ref() == active_path.as_ref() + }) + .cloned(); + if recommended.is_some() { + return recommended; + } + } + + // Then try the first PythonEnv with ipykernel matching the language + let python_env = self + .kernel_specifications_for_worktree(worktree_id) + .find(|spec| { + matches!(spec, KernelSpecification::PythonEnv(_)) + && spec.has_ipykernel() + && spec.language().as_ref().to_lowercase() == language_name }) - } else { - selected_kernelspec + .cloned(); + if python_env.is_some() { + return python_env; } + + // Fall back to legacy name-based and language-based matching + self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx) } fn kernelspec_legacy_by_lang_only( @@ -244,7 +311,6 @@ impl ReplStore { if let (Some(selected), KernelSpecification::Jupyter(runtime_specification)) = (selected_kernel, runtime_specification) { - // Top priority is the selected kernel return runtime_specification.name.to_lowercase() == selected.to_lowercase(); } false @@ -255,20 +321,10 @@ impl ReplStore { return Some(found_by_name); } + let language_name = language_at_cursor.code_fence_block_name().to_lowercase(); self.kernel_specifications_for_worktree(worktree_id) - .find(|kernel_option| match kernel_option { - KernelSpecification::Jupyter(runtime_specification) => { - runtime_specification.kernelspec.language.to_lowercase() - == language_at_cursor.code_fence_block_name().to_lowercase() - } - KernelSpecification::PythonEnv(runtime_specification) => { - runtime_specification.kernelspec.language.to_lowercase() - == language_at_cursor.code_fence_block_name().to_lowercase() - } - KernelSpecification::Remote(remote_spec) => { - remote_spec.kernelspec.language.to_lowercase() - == language_at_cursor.code_fence_block_name().to_lowercase() - } + .find(|spec| { + spec.has_ipykernel() && spec.language().as_ref().to_lowercase() == language_name }) .cloned() } diff --git a/crates/repl/src/session.rs b/crates/repl/src/session.rs index 2fda4f1bedeba54a9cc07146d5f3b81706073948..fcb06c1409c00a6eebf25d48fde89d63ea1d070e 100644 --- a/crates/repl/src/session.rs +++ b/crates/repl/src/session.rs @@ -277,8 +277,7 @@ impl Session { let session_view = cx.entity(); let kernel = match self.kernel_specification.clone() { - KernelSpecification::Jupyter(kernel_specification) - | KernelSpecification::PythonEnv(kernel_specification) => NativeRunningKernel::new( + KernelSpecification::Jupyter(kernel_specification) => NativeRunningKernel::new( kernel_specification, entity_id, working_directory, @@ -287,6 +286,15 @@ impl Session { window, cx, ), + KernelSpecification::PythonEnv(env_specification) => NativeRunningKernel::new( + env_specification.as_local_spec(), + entity_id, + working_directory, + self.fs.clone(), + session_view, + window, + cx, + ), KernelSpecification::Remote(remote_kernel_specification) => RemoteRunningKernel::new( remote_kernel_specification, working_directory, diff --git a/crates/zed/src/zed/quick_action_bar/repl_menu.rs b/crates/zed/src/zed/quick_action_bar/repl_menu.rs index 1ebdf35bb93824b7881afabe289a07feb93f8135..45c0dd75f17a155d74190f4bcbbcf5296cebacdb 100644 --- a/crates/zed/src/zed/quick_action_bar/repl_menu.rs +++ b/crates/zed/src/zed/quick_action_bar/repl_menu.rs @@ -283,6 +283,21 @@ impl QuickActionBar { return div().into_any_element(); }; + let store = repl::ReplStore::global(cx); + if !store.read(cx).has_python_kernelspecs(worktree_id) { + if let Some(project) = editor + .read(cx) + .workspace() + .map(|workspace| workspace.read(cx).project().clone()) + { + store + .update(cx, |store, cx| { + store.refresh_python_kernelspecs(worktree_id, &project, cx) + }) + .detach_and_log_err(cx); + } + } + let session = repl::session(editor.downgrade(), cx); let current_kernelspec = match session { @@ -301,7 +316,17 @@ impl QuickActionBar { KernelSelector::new( { Box::new(move |kernelspec, window, cx| { - repl::assign_kernelspec(kernelspec, editor.downgrade(), window, cx).ok(); + if kernelspec.has_ipykernel() { + repl::assign_kernelspec(kernelspec, editor.downgrade(), window, cx).ok(); + } else { + repl::install_ipykernel_and_assign( + kernelspec, + editor.downgrade(), + window, + cx, + ) + .ok(); + } }) }, worktree_id,