490 lines · 17.4 KB · Rust
RawDownload
use crate::KERNEL_DOCS_URL;use crate::kernels::KernelSpecification;use crate::repl_store::ReplStore; use gpui::{AnyView, DismissEvent, FontWeight, SharedString, Task};use picker::{Picker, PickerDelegate};use project::WorktreeId;use std::sync::Arc;use ui::{ListItem, ListItemSpacing, PopoverMenu, PopoverMenuHandle, PopoverTrigger, prelude::*}; type OnSelect = Box<dyn Fn(KernelSpecification, &mut Window, &mut App)>; #[derive(Clone)]pub enum KernelPickerEntry {    SectionHeader(SharedString),    Kernel {        spec: KernelSpecification,        is_recommended: bool,    },} fn build_grouped_entries(store: &ReplStore, worktree_id: WorktreeId) -> Vec<KernelPickerEntry> {    let mut entries = Vec::new();    let mut recommended_entry: Option<KernelPickerEntry> = None;    let mut found_selected = false;    let selected_kernel = store.selected_kernel(worktree_id);     let mut python_envs = Vec::new();    let mut jupyter_kernels = Vec::new();    let mut wsl_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);        let is_selected = selected_kernel.map_or(false, |s| s == spec);         if is_selected {            recommended_entry = Some(KernelPickerEntry::Kernel {                spec: spec.clone(),                is_recommended: true,            });            found_selected = true;        } else if is_recommended && !found_selected {            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::JupyterServer(_) | KernelSpecification::SshRemote(_) => {                remote_kernels.push(KernelPickerEntry::Kernel {                    spec: spec.clone(),                    is_recommended,                });            }            KernelSpecification::WslRemote(_) => {                wsl_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);    }     // WSL Kernels section    if !wsl_kernels.is_empty() {        entries.push(KernelPickerEntry::SectionHeader("WSL Kernels".into()));        entries.extend(wsl_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<T, TT>where    T: PopoverTrigger + ButtonCommon,    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,{    handle: Option<PopoverMenuHandle<Picker<KernelPickerDelegate>>>,    on_select: OnSelect,    trigger: T,    tooltip: TT,    info_text: Option<SharedString>,    worktree_id: WorktreeId,} pub struct KernelPickerDelegate {    all_entries: Vec<KernelPickerEntry>,    filtered_entries: Vec<KernelPickerEntry>,    selected_kernelspec: Option<KernelSpecification>,    selected_index: usize,    on_select: OnSelect,} impl<T, TT> KernelSelector<T, TT>where    T: PopoverTrigger + ButtonCommon,    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,{    pub fn new(on_select: OnSelect, worktree_id: WorktreeId, trigger: T, tooltip: TT) -> Self {        KernelSelector {            on_select,            handle: None,            trigger,            tooltip,            info_text: None,            worktree_id,        }    }     pub fn with_handle(mut self, handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>) -> Self {        self.handle = Some(handle);        self    }     pub fn with_info_text(mut self, text: impl Into<SharedString>) -> Self {        self.info_text = Some(text.into());        self    }} 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_entries.len()    }     fn selected_index(&self) -> usize {        self.selected_index    }     fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context<Picker<Self>>) {        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();    }     fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {        "Select a kernel...".into()    }     fn update_matches(        &mut self,        query: String,        _window: &mut Window,        _cx: &mut Context<Picker<Self>>,    ) -> Task<()> {        if query.is_empty() {            self.filtered_entries = self.all_entries.clone();        } else {            let query_lower = query.to_lowercase();            let mut filtered = Vec::new();            let mut pending_header: Option<KernelPickerEntry> = 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.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<usize> {        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<Picker<Self>>) {        if let Some(KernelPickerEntry::Kernel { spec, .. }) =            self.filtered_entries.get(self.selected_index)        {            (self.on_select)(spec.clone(), window, cx);            cx.emit(DismissEvent);        }    }     fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context<Picker<Self>>) {}     fn render_match(        &self,        ix: usize,        selected: bool,        _: &mut Window,        cx: &mut Context<Picker<Self>>,    ) -> Option<Self::ListItem> {        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),                    ),            ),            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::WslRemote(_) => Some(spec.path().to_string()),                    KernelSpecification::PythonEnv(_)                    | KernelSpecification::JupyterServer(_)                    | KernelSpecification::SshRemote(_) => {                        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(                            h_flex()                                .w_full()                                .gap_3()                                .when(!has_ipykernel, |flex| flex.opacity(0.5))                                .child(icon.color(Color::Default).size(IconSize::Medium))                                .child(                                    v_flex()                                        .flex_grow()                                        .overflow_x_hidden()                                        .gap_0p5()                                        .child(                                            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(subtitle, |flex, subtitle| {                                            flex.child(                                                div().overflow_x_hidden().text_ellipsis().child(                                                    Label::new(subtitle)                                                        .size(LabelSize::Small)                                                        .color(Color::Muted),                                                ),                                            )                                        }),                                ),                        )                        .when(is_currently_selected, |item| {                            item.end_slot(                                Icon::new(IconName::Check)                                    .color(Color::Accent)                                    .size(IconSize::Small),                            )                        }),                )            }        }    }     fn render_footer(        &self,        _: &mut Window,        cx: &mut Context<Picker<Self>>,    ) -> Option<gpui::AnyElement> {        Some(            h_flex()                .w_full()                .border_t_1()                .border_color(cx.theme().colors().border_variant)                .p_1()                .gap_4()                .child(                    Button::new("kernel-docs", "Kernel Docs")                        .end_icon(                            Icon::new(IconName::ArrowUpRight)                                .size(IconSize::Small)                                .color(Color::Muted),                        )                        .on_click(move |_, _, cx| cx.open_url(KERNEL_DOCS_URL)),                )                .into_any(),        )    }} impl<T, TT> RenderOnce for KernelSelector<T, TT>where    T: PopoverTrigger + ButtonCommon,    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,{    fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {        let store = ReplStore::global(cx);        store.update(cx, |store, cx| store.ensure_kernelspecs(cx));        let store = store.read(cx);         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_entries: all_entries.clone(),            filtered_entries: all_entries,            selected_kernelspec,            selected_index,        };         let picker_view = cx.new(|cx| {            Picker::list(delegate, window, cx)                .list_measure_all()                .width(rems(34.))                .max_height(Some(rems(24.).into()))        });         PopoverMenu::new("kernel-switcher")            .menu(move |_window, _cx| Some(picker_view.clone()))            .trigger_with_tooltip(self.trigger, self.tooltip)            .attach(gpui::Corner::BottomLeft)            .when_some(self.handle, |menu, handle| menu.with_handle(handle))    }}