1use crate::KERNEL_DOCS_URL;
2use crate::kernels::KernelSpecification;
3use crate::repl_store::ReplStore;
4
5use gpui::{AnyView, DismissEvent, FontWeight, SharedString, Task};
6use picker::{Picker, PickerDelegate};
7use project::WorktreeId;
8use std::sync::Arc;
9use ui::{ListItem, ListItemSpacing, PopoverMenu, PopoverMenuHandle, PopoverTrigger, prelude::*};
10
11type OnSelect = Box<dyn Fn(KernelSpecification, &mut Window, &mut App)>;
12
13#[derive(Clone)]
14pub enum KernelPickerEntry {
15 SectionHeader(SharedString),
16 Kernel {
17 spec: KernelSpecification,
18 is_recommended: bool,
19 },
20}
21
22fn build_grouped_entries(store: &ReplStore, worktree_id: WorktreeId) -> Vec<KernelPickerEntry> {
23 let mut entries = Vec::new();
24 let mut recommended_entry: Option<KernelPickerEntry> = None;
25 let mut found_selected = false;
26 let selected_kernel = store.selected_kernel(worktree_id);
27
28 let mut python_envs = Vec::new();
29 let mut jupyter_kernels = Vec::new();
30 let mut remote_kernels = Vec::new();
31
32 for spec in store.kernel_specifications_for_worktree(worktree_id) {
33 let is_recommended = store.is_recommended_kernel(worktree_id, spec);
34 let is_selected = selected_kernel.map_or(false, |s| s == spec);
35
36 if is_selected {
37 recommended_entry = Some(KernelPickerEntry::Kernel {
38 spec: spec.clone(),
39 is_recommended: true,
40 });
41 found_selected = true;
42 } else if is_recommended && !found_selected {
43 recommended_entry = Some(KernelPickerEntry::Kernel {
44 spec: spec.clone(),
45 is_recommended: true,
46 });
47 }
48
49 match spec {
50 KernelSpecification::PythonEnv(_) => {
51 python_envs.push(KernelPickerEntry::Kernel {
52 spec: spec.clone(),
53 is_recommended,
54 });
55 }
56 KernelSpecification::Jupyter(_) => {
57 jupyter_kernels.push(KernelPickerEntry::Kernel {
58 spec: spec.clone(),
59 is_recommended,
60 });
61 }
62 KernelSpecification::JupyterServer(_)
63 | KernelSpecification::SshRemote(_)
64 | KernelSpecification::WslRemote(_) => {
65 remote_kernels.push(KernelPickerEntry::Kernel {
66 spec: spec.clone(),
67 is_recommended,
68 });
69 }
70 }
71 }
72
73 // Sort Python envs: has_ipykernel first, then by name
74 python_envs.sort_by(|a, b| {
75 let (spec_a, spec_b) = match (a, b) {
76 (
77 KernelPickerEntry::Kernel { spec: sa, .. },
78 KernelPickerEntry::Kernel { spec: sb, .. },
79 ) => (sa, sb),
80 _ => return std::cmp::Ordering::Equal,
81 };
82 spec_b
83 .has_ipykernel()
84 .cmp(&spec_a.has_ipykernel())
85 .then_with(|| spec_a.name().cmp(&spec_b.name()))
86 });
87
88 // Recommended section
89 if let Some(rec) = recommended_entry {
90 entries.push(KernelPickerEntry::SectionHeader("Recommended".into()));
91 entries.push(rec);
92 }
93
94 // Python Environments section
95 if !python_envs.is_empty() {
96 entries.push(KernelPickerEntry::SectionHeader(
97 "Python Environments".into(),
98 ));
99 entries.extend(python_envs);
100 }
101
102 // Jupyter Kernels section
103 if !jupyter_kernels.is_empty() {
104 entries.push(KernelPickerEntry::SectionHeader("Jupyter Kernels".into()));
105 entries.extend(jupyter_kernels);
106 }
107
108 // Remote section
109 if !remote_kernels.is_empty() {
110 entries.push(KernelPickerEntry::SectionHeader("Remote Servers".into()));
111 entries.extend(remote_kernels);
112 }
113
114 entries
115}
116
117#[derive(IntoElement)]
118pub struct KernelSelector<T, TT>
119where
120 T: PopoverTrigger + ButtonCommon,
121 TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
122{
123 handle: Option<PopoverMenuHandle<Picker<KernelPickerDelegate>>>,
124 on_select: OnSelect,
125 trigger: T,
126 tooltip: TT,
127 info_text: Option<SharedString>,
128 worktree_id: WorktreeId,
129}
130
131pub struct KernelPickerDelegate {
132 all_entries: Vec<KernelPickerEntry>,
133 filtered_entries: Vec<KernelPickerEntry>,
134 selected_kernelspec: Option<KernelSpecification>,
135 selected_index: usize,
136 on_select: OnSelect,
137}
138
139impl<T, TT> KernelSelector<T, TT>
140where
141 T: PopoverTrigger + ButtonCommon,
142 TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
143{
144 pub fn new(on_select: OnSelect, worktree_id: WorktreeId, trigger: T, tooltip: TT) -> Self {
145 KernelSelector {
146 on_select,
147 handle: None,
148 trigger,
149 tooltip,
150 info_text: None,
151 worktree_id,
152 }
153 }
154
155 pub fn with_handle(mut self, handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>) -> Self {
156 self.handle = Some(handle);
157 self
158 }
159
160 pub fn with_info_text(mut self, text: impl Into<SharedString>) -> Self {
161 self.info_text = Some(text.into());
162 self
163 }
164}
165
166impl KernelPickerDelegate {
167 fn first_selectable_index(entries: &[KernelPickerEntry]) -> usize {
168 entries
169 .iter()
170 .position(|e| matches!(e, KernelPickerEntry::Kernel { .. }))
171 .unwrap_or(0)
172 }
173
174 fn next_selectable_index(&self, from: usize, direction: i32) -> usize {
175 let len = self.filtered_entries.len();
176 if len == 0 {
177 return 0;
178 }
179
180 let mut index = from as i32 + direction;
181 while index >= 0 && (index as usize) < len {
182 if matches!(
183 self.filtered_entries.get(index as usize),
184 Some(KernelPickerEntry::Kernel { .. })
185 ) {
186 return index as usize;
187 }
188 index += direction;
189 }
190
191 from
192 }
193}
194
195impl PickerDelegate for KernelPickerDelegate {
196 type ListItem = ListItem;
197
198 fn match_count(&self) -> usize {
199 self.filtered_entries.len()
200 }
201
202 fn selected_index(&self) -> usize {
203 self.selected_index
204 }
205
206 fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context<Picker<Self>>) {
207 if matches!(
208 self.filtered_entries.get(ix),
209 Some(KernelPickerEntry::SectionHeader(_))
210 ) {
211 let forward = self.next_selectable_index(ix, 1);
212 if forward != ix {
213 self.selected_index = forward;
214 } else {
215 self.selected_index = self.next_selectable_index(ix, -1);
216 }
217 } else {
218 self.selected_index = ix;
219 }
220
221 if let Some(KernelPickerEntry::Kernel { spec, .. }) =
222 self.filtered_entries.get(self.selected_index)
223 {
224 self.selected_kernelspec = Some(spec.clone());
225 }
226 cx.notify();
227 }
228
229 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
230 "Select a kernel...".into()
231 }
232
233 fn update_matches(
234 &mut self,
235 query: String,
236 _window: &mut Window,
237 _cx: &mut Context<Picker<Self>>,
238 ) -> Task<()> {
239 if query.is_empty() {
240 self.filtered_entries = self.all_entries.clone();
241 } else {
242 let query_lower = query.to_lowercase();
243 let mut filtered = Vec::new();
244 let mut pending_header: Option<KernelPickerEntry> = None;
245
246 for entry in &self.all_entries {
247 match entry {
248 KernelPickerEntry::SectionHeader(_) => {
249 pending_header = Some(entry.clone());
250 }
251 KernelPickerEntry::Kernel { spec, .. } => {
252 if spec.name().to_lowercase().contains(&query_lower) {
253 if let Some(header) = pending_header.take() {
254 filtered.push(header);
255 }
256 filtered.push(entry.clone());
257 }
258 }
259 }
260 }
261
262 self.filtered_entries = filtered;
263 }
264
265 self.selected_index = Self::first_selectable_index(&self.filtered_entries);
266 if let Some(KernelPickerEntry::Kernel { spec, .. }) =
267 self.filtered_entries.get(self.selected_index)
268 {
269 self.selected_kernelspec = Some(spec.clone());
270 }
271
272 Task::ready(())
273 }
274
275 fn separators_after_indices(&self) -> Vec<usize> {
276 let mut separators = Vec::new();
277 for (index, entry) in self.filtered_entries.iter().enumerate() {
278 if matches!(entry, KernelPickerEntry::SectionHeader(_)) && index > 0 {
279 separators.push(index - 1);
280 }
281 }
282 separators
283 }
284
285 fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
286 if let Some(KernelPickerEntry::Kernel { spec, .. }) =
287 self.filtered_entries.get(self.selected_index)
288 {
289 (self.on_select)(spec.clone(), window, cx);
290 cx.emit(DismissEvent);
291 }
292 }
293
294 fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context<Picker<Self>>) {}
295
296 fn render_match(
297 &self,
298 ix: usize,
299 selected: bool,
300 _: &mut Window,
301 cx: &mut Context<Picker<Self>>,
302 ) -> Option<Self::ListItem> {
303 let entry = self.filtered_entries.get(ix)?;
304
305 match entry {
306 KernelPickerEntry::SectionHeader(title) => Some(
307 ListItem::new(ix)
308 .inset(true)
309 .spacing(ListItemSpacing::Dense)
310 .selectable(false)
311 .child(
312 Label::new(title.clone())
313 .size(LabelSize::Small)
314 .weight(FontWeight::SEMIBOLD)
315 .color(Color::Muted),
316 ),
317 ),
318 KernelPickerEntry::Kernel {
319 spec,
320 is_recommended,
321 } => {
322 let is_currently_selected = self.selected_kernelspec.as_ref() == Some(spec);
323 let icon = spec.icon(cx);
324 let has_ipykernel = spec.has_ipykernel();
325
326 let subtitle = match spec {
327 KernelSpecification::Jupyter(_) => None,
328 KernelSpecification::PythonEnv(_)
329 | KernelSpecification::JupyterServer(_)
330 | KernelSpecification::SshRemote(_)
331 | KernelSpecification::WslRemote(_) => {
332 let env_kind = spec.environment_kind_label();
333 let path = spec.path();
334 match env_kind {
335 Some(kind) => Some(format!("{} \u{2013} {}", kind, path)),
336 None => Some(path.to_string()),
337 }
338 }
339 };
340
341 Some(
342 ListItem::new(ix)
343 .inset(true)
344 .spacing(ListItemSpacing::Sparse)
345 .toggle_state(selected)
346 .child(
347 h_flex()
348 .w_full()
349 .gap_3()
350 .when(!has_ipykernel, |flex| flex.opacity(0.5))
351 .child(icon.color(Color::Default).size(IconSize::Medium))
352 .child(
353 v_flex()
354 .flex_grow()
355 .overflow_x_hidden()
356 .gap_0p5()
357 .child(
358 h_flex()
359 .gap_1()
360 .child(
361 div()
362 .overflow_x_hidden()
363 .flex_shrink()
364 .text_ellipsis()
365 .child(
366 Label::new(spec.name())
367 .weight(FontWeight::MEDIUM)
368 .size(LabelSize::Default),
369 ),
370 )
371 .when(*is_recommended, |flex| {
372 flex.child(
373 Label::new("Recommended")
374 .size(LabelSize::XSmall)
375 .color(Color::Accent),
376 )
377 })
378 .when(!has_ipykernel, |flex| {
379 flex.child(
380 Label::new("ipykernel not installed")
381 .size(LabelSize::XSmall)
382 .color(Color::Warning),
383 )
384 }),
385 )
386 .when_some(subtitle, |flex, subtitle| {
387 flex.child(
388 div().overflow_x_hidden().text_ellipsis().child(
389 Label::new(subtitle)
390 .size(LabelSize::Small)
391 .color(Color::Muted),
392 ),
393 )
394 }),
395 ),
396 )
397 .when(is_currently_selected, |item| {
398 item.end_slot(
399 Icon::new(IconName::Check)
400 .color(Color::Accent)
401 .size(IconSize::Small),
402 )
403 }),
404 )
405 }
406 }
407 }
408
409 fn render_footer(
410 &self,
411 _: &mut Window,
412 cx: &mut Context<Picker<Self>>,
413 ) -> Option<gpui::AnyElement> {
414 Some(
415 h_flex()
416 .w_full()
417 .border_t_1()
418 .border_color(cx.theme().colors().border_variant)
419 .p_1()
420 .gap_4()
421 .child(
422 Button::new("kernel-docs", "Kernel Docs")
423 .icon(IconName::ArrowUpRight)
424 .icon_size(IconSize::Small)
425 .icon_color(Color::Muted)
426 .icon_position(IconPosition::End)
427 .on_click(move |_, _, cx| cx.open_url(KERNEL_DOCS_URL)),
428 )
429 .into_any(),
430 )
431 }
432}
433
434impl<T, TT> RenderOnce for KernelSelector<T, TT>
435where
436 T: PopoverTrigger + ButtonCommon,
437 TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
438{
439 fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
440 let store = ReplStore::global(cx).read(cx);
441
442 let all_entries = build_grouped_entries(store, self.worktree_id);
443 let selected_kernelspec = store.active_kernelspec(self.worktree_id, None, cx);
444 let selected_index = all_entries
445 .iter()
446 .position(|entry| {
447 if let KernelPickerEntry::Kernel { spec, .. } = entry {
448 selected_kernelspec.as_ref() == Some(spec)
449 } else {
450 false
451 }
452 })
453 .unwrap_or_else(|| KernelPickerDelegate::first_selectable_index(&all_entries));
454
455 let delegate = KernelPickerDelegate {
456 on_select: self.on_select,
457 all_entries: all_entries.clone(),
458 filtered_entries: all_entries,
459 selected_kernelspec,
460 selected_index,
461 };
462
463 let picker_view = cx.new(|cx| {
464 Picker::list(delegate, window, cx)
465 .list_measure_all()
466 .width(rems(34.))
467 .max_height(Some(rems(24.).into()))
468 });
469
470 PopoverMenu::new("kernel-switcher")
471 .menu(move |_window, _cx| Some(picker_view.clone()))
472 .trigger_with_tooltip(self.trigger, self.tooltip)
473 .attach(gpui::Corner::BottomLeft)
474 .when_some(self.handle, |menu, handle| menu.with_handle(handle))
475 }
476}