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