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