notebook_ui.rs

  1#![allow(unused, dead_code)]
  2use std::future::Future;
  3use std::{path::PathBuf, sync::Arc};
  4
  5use anyhow::{Context as _, Result};
  6use client::proto::ViewId;
  7use collections::HashMap;
  8use feature_flags::{FeatureFlagAppExt as _, NotebookFeatureFlag};
  9use futures::future::Shared;
 10use futures::FutureExt;
 11use gpui::{
 12    actions, list, prelude::*, AnyElement, AppContext, EventEmitter, FocusHandle, FocusableView,
 13    ListScrollEvent, ListState, Model, Point, Task, View,
 14};
 15use language::{Language, LanguageRegistry};
 16use project::{Project, ProjectEntryId, ProjectPath};
 17use ui::{prelude::*, Tooltip};
 18use workspace::item::{ItemEvent, TabContentParams};
 19use workspace::searchable::SearchableItemHandle;
 20use workspace::{Item, ItemHandle, ProjectItem, ToolbarItemLocation};
 21use workspace::{ToolbarItemEvent, ToolbarItemView};
 22
 23use super::{Cell, CellPosition, RenderableCell};
 24
 25use nbformat::v4::CellId;
 26use nbformat::v4::Metadata as NotebookMetadata;
 27
 28actions!(
 29    notebook,
 30    [
 31        OpenNotebook,
 32        RunAll,
 33        ClearOutputs,
 34        MoveCellUp,
 35        MoveCellDown,
 36        AddMarkdownBlock,
 37        AddCodeBlock,
 38    ]
 39);
 40
 41pub(crate) const MAX_TEXT_BLOCK_WIDTH: f32 = 9999.0;
 42pub(crate) const SMALL_SPACING_SIZE: f32 = 8.0;
 43pub(crate) const MEDIUM_SPACING_SIZE: f32 = 12.0;
 44pub(crate) const LARGE_SPACING_SIZE: f32 = 16.0;
 45pub(crate) const GUTTER_WIDTH: f32 = 19.0;
 46pub(crate) const CODE_BLOCK_INSET: f32 = MEDIUM_SPACING_SIZE;
 47pub(crate) const CONTROL_SIZE: f32 = 20.0;
 48
 49pub fn init(cx: &mut AppContext) {
 50    if cx.has_flag::<NotebookFeatureFlag>() || std::env::var("LOCAL_NOTEBOOK_DEV").is_ok() {
 51        workspace::register_project_item::<NotebookEditor>(cx);
 52    }
 53
 54    cx.observe_flag::<NotebookFeatureFlag, _>({
 55        move |is_enabled, cx| {
 56            if is_enabled {
 57                workspace::register_project_item::<NotebookEditor>(cx);
 58            } else {
 59                // todo: there is no way to unregister a project item, so if the feature flag
 60                // gets turned off they need to restart Zed.
 61            }
 62        }
 63    })
 64    .detach();
 65}
 66
 67pub struct NotebookEditor {
 68    languages: Arc<LanguageRegistry>,
 69    project: Model<Project>,
 70
 71    focus_handle: FocusHandle,
 72    notebook_item: Model<NotebookItem>,
 73
 74    remote_id: Option<ViewId>,
 75    cell_list: ListState,
 76
 77    selected_cell_index: usize,
 78    cell_order: Vec<CellId>,
 79    cell_map: HashMap<CellId, Cell>,
 80}
 81
 82impl NotebookEditor {
 83    pub fn new(
 84        project: Model<Project>,
 85        notebook_item: Model<NotebookItem>,
 86        cx: &mut ViewContext<Self>,
 87    ) -> Self {
 88        let focus_handle = cx.focus_handle();
 89
 90        let languages = project.read(cx).languages().clone();
 91        let language_name = notebook_item.read(cx).language_name();
 92
 93        let notebook_language = notebook_item.read(cx).notebook_language();
 94        let notebook_language = cx.spawn(|_, _| notebook_language).shared();
 95
 96        let mut cell_order = vec![]; // Vec<CellId>
 97        let mut cell_map = HashMap::default(); // HashMap<CellId, Cell>
 98
 99        for (index, cell) in notebook_item
100            .read(cx)
101            .notebook
102            .clone()
103            .cells
104            .iter()
105            .enumerate()
106        {
107            let cell_id = cell.id();
108            cell_order.push(cell_id.clone());
109            cell_map.insert(
110                cell_id.clone(),
111                Cell::load(cell, &languages, notebook_language.clone(), cx),
112            );
113        }
114
115        let view = cx.view().downgrade();
116        let cell_count = cell_order.len();
117
118        let this = cx.view();
119        let cell_list = ListState::new(
120            cell_count,
121            gpui::ListAlignment::Top,
122            px(1000.),
123            move |ix, cx| {
124                view.upgrade()
125                    .and_then(|notebook_handle| {
126                        notebook_handle.update(cx, |notebook, cx| {
127                            notebook
128                                .cell_order
129                                .get(ix)
130                                .and_then(|cell_id| notebook.cell_map.get(cell_id))
131                                .map(|cell| notebook.render_cell(ix, cell, cx).into_any_element())
132                        })
133                    })
134                    .unwrap_or_else(|| div().into_any())
135            },
136        );
137
138        Self {
139            project,
140            languages: languages.clone(),
141            focus_handle,
142            notebook_item,
143            remote_id: None,
144            cell_list,
145            selected_cell_index: 0,
146            cell_order: cell_order.clone(),
147            cell_map: cell_map.clone(),
148        }
149    }
150
151    fn has_outputs(&self, cx: &ViewContext<Self>) -> bool {
152        self.cell_map.values().any(|cell| {
153            if let Cell::Code(code_cell) = cell {
154                code_cell.read(cx).has_outputs()
155            } else {
156                false
157            }
158        })
159    }
160
161    fn is_dirty(&self, cx: &AppContext) -> bool {
162        self.cell_map.values().any(|cell| {
163            if let Cell::Code(code_cell) = cell {
164                code_cell.read(cx).is_dirty(cx)
165            } else {
166                false
167            }
168        })
169    }
170
171    fn clear_outputs(&mut self, cx: &mut ViewContext<Self>) {
172        for cell in self.cell_map.values() {
173            if let Cell::Code(code_cell) = cell {
174                code_cell.update(cx, |cell, _cx| {
175                    cell.clear_outputs();
176                });
177            }
178        }
179    }
180
181    fn run_cells(&mut self, cx: &mut ViewContext<Self>) {
182        println!("Cells would all run here, if that was implemented!");
183    }
184
185    fn open_notebook(&mut self, _: &OpenNotebook, _cx: &mut ViewContext<Self>) {
186        println!("Open notebook triggered");
187    }
188
189    fn move_cell_up(&mut self, cx: &mut ViewContext<Self>) {
190        println!("Move cell up triggered");
191    }
192
193    fn move_cell_down(&mut self, cx: &mut ViewContext<Self>) {
194        println!("Move cell down triggered");
195    }
196
197    fn add_markdown_block(&mut self, cx: &mut ViewContext<Self>) {
198        println!("Add markdown block triggered");
199    }
200
201    fn add_code_block(&mut self, cx: &mut ViewContext<Self>) {
202        println!("Add code block triggered");
203    }
204
205    fn cell_count(&self) -> usize {
206        self.cell_map.len()
207    }
208
209    fn selected_index(&self) -> usize {
210        self.selected_cell_index
211    }
212
213    pub fn set_selected_index(
214        &mut self,
215        index: usize,
216        jump_to_index: bool,
217        cx: &mut ViewContext<Self>,
218    ) {
219        // let previous_index = self.selected_cell_index;
220        self.selected_cell_index = index;
221        let current_index = self.selected_cell_index;
222
223        // in the future we may have some `on_cell_change` event that we want to fire here
224
225        if jump_to_index {
226            self.jump_to_cell(current_index, cx);
227        }
228    }
229
230    pub fn select_next(&mut self, _: &menu::SelectNext, cx: &mut ViewContext<Self>) {
231        let count = self.cell_count();
232        if count > 0 {
233            let index = self.selected_index();
234            let ix = if index == count - 1 {
235                count - 1
236            } else {
237                index + 1
238            };
239            self.set_selected_index(ix, true, cx);
240            cx.notify();
241        }
242    }
243
244    pub fn select_previous(&mut self, _: &menu::SelectPrev, cx: &mut ViewContext<Self>) {
245        let count = self.cell_count();
246        if count > 0 {
247            let index = self.selected_index();
248            let ix = if index == 0 { 0 } else { index - 1 };
249            self.set_selected_index(ix, true, cx);
250            cx.notify();
251        }
252    }
253
254    pub fn select_first(&mut self, _: &menu::SelectFirst, cx: &mut ViewContext<Self>) {
255        let count = self.cell_count();
256        if count > 0 {
257            self.set_selected_index(0, true, cx);
258            cx.notify();
259        }
260    }
261
262    pub fn select_last(&mut self, _: &menu::SelectLast, cx: &mut ViewContext<Self>) {
263        let count = self.cell_count();
264        if count > 0 {
265            self.set_selected_index(count - 1, true, cx);
266            cx.notify();
267        }
268    }
269
270    fn jump_to_cell(&mut self, index: usize, _cx: &mut ViewContext<Self>) {
271        self.cell_list.scroll_to_reveal_item(index);
272    }
273
274    fn button_group(cx: &ViewContext<Self>) -> Div {
275        v_flex()
276            .gap(Spacing::Small.rems(cx))
277            .items_center()
278            .w(px(CONTROL_SIZE + 4.0))
279            .overflow_hidden()
280            .rounded(px(5.))
281            .bg(cx.theme().colors().title_bar_background)
282            .p_px()
283            .border_1()
284            .border_color(cx.theme().colors().border)
285    }
286
287    fn render_notebook_control(
288        id: impl Into<SharedString>,
289        icon: IconName,
290        _cx: &ViewContext<Self>,
291    ) -> IconButton {
292        let id: ElementId = ElementId::Name(id.into());
293        IconButton::new(id, icon).width(px(CONTROL_SIZE).into())
294    }
295
296    fn render_notebook_controls(&self, cx: &ViewContext<Self>) -> impl IntoElement {
297        let has_outputs = self.has_outputs(cx);
298
299        v_flex()
300            .max_w(px(CONTROL_SIZE + 4.0))
301            .items_center()
302            .gap(Spacing::XXLarge.rems(cx))
303            .justify_between()
304            .flex_none()
305            .h_full()
306            .py(Spacing::XLarge.px(cx))
307            .child(
308                v_flex()
309                    .gap(Spacing::Large.rems(cx))
310                    .child(
311                        Self::button_group(cx)
312                            .child(
313                                Self::render_notebook_control("run-all-cells", IconName::Play, cx)
314                                    .tooltip(move |cx| {
315                                        Tooltip::for_action("Execute all cells", &RunAll, cx)
316                                    })
317                                    .on_click(|_, cx| {
318                                        cx.dispatch_action(Box::new(RunAll));
319                                    }),
320                            )
321                            .child(
322                                Self::render_notebook_control(
323                                    "clear-all-outputs",
324                                    IconName::ListX,
325                                    cx,
326                                )
327                                .disabled(!has_outputs)
328                                .tooltip(move |cx| {
329                                    Tooltip::for_action("Clear all outputs", &ClearOutputs, cx)
330                                })
331                                .on_click(|_, cx| {
332                                    cx.dispatch_action(Box::new(ClearOutputs));
333                                }),
334                            ),
335                    )
336                    .child(
337                        Self::button_group(cx)
338                            .child(
339                                Self::render_notebook_control(
340                                    "move-cell-up",
341                                    IconName::ArrowUp,
342                                    cx,
343                                )
344                                .tooltip(move |cx| {
345                                    Tooltip::for_action("Move cell up", &MoveCellUp, cx)
346                                })
347                                .on_click(|_, cx| {
348                                    cx.dispatch_action(Box::new(MoveCellUp));
349                                }),
350                            )
351                            .child(
352                                Self::render_notebook_control(
353                                    "move-cell-down",
354                                    IconName::ArrowDown,
355                                    cx,
356                                )
357                                .tooltip(move |cx| {
358                                    Tooltip::for_action("Move cell down", &MoveCellDown, cx)
359                                })
360                                .on_click(|_, cx| {
361                                    cx.dispatch_action(Box::new(MoveCellDown));
362                                }),
363                            ),
364                    )
365                    .child(
366                        Self::button_group(cx)
367                            .child(
368                                Self::render_notebook_control(
369                                    "new-markdown-cell",
370                                    IconName::Plus,
371                                    cx,
372                                )
373                                .tooltip(move |cx| {
374                                    Tooltip::for_action("Add markdown block", &AddMarkdownBlock, cx)
375                                })
376                                .on_click(|_, cx| {
377                                    cx.dispatch_action(Box::new(AddMarkdownBlock));
378                                }),
379                            )
380                            .child(
381                                Self::render_notebook_control("new-code-cell", IconName::Code, cx)
382                                    .tooltip(move |cx| {
383                                        Tooltip::for_action("Add code block", &AddCodeBlock, cx)
384                                    })
385                                    .on_click(|_, cx| {
386                                        cx.dispatch_action(Box::new(AddCodeBlock));
387                                    }),
388                            ),
389                    ),
390            )
391            .child(
392                v_flex()
393                    .gap(Spacing::Large.rems(cx))
394                    .items_center()
395                    .child(Self::render_notebook_control(
396                        "more-menu",
397                        IconName::Ellipsis,
398                        cx,
399                    ))
400                    .child(
401                        Self::button_group(cx)
402                            .child(IconButton::new("repl", IconName::ReplNeutral)),
403                    ),
404            )
405    }
406
407    fn cell_position(&self, index: usize) -> CellPosition {
408        match index {
409            0 => CellPosition::First,
410            index if index == self.cell_count() - 1 => CellPosition::Last,
411            _ => CellPosition::Middle,
412        }
413    }
414
415    fn render_cell(
416        &self,
417        index: usize,
418        cell: &Cell,
419        cx: &mut ViewContext<Self>,
420    ) -> impl IntoElement {
421        let cell_position = self.cell_position(index);
422
423        let is_selected = index == self.selected_cell_index;
424
425        match cell {
426            Cell::Code(cell) => {
427                cell.update(cx, |cell, _cx| {
428                    cell.set_selected(is_selected)
429                        .set_cell_position(cell_position);
430                });
431                cell.clone().into_any_element()
432            }
433            Cell::Markdown(cell) => {
434                cell.update(cx, |cell, _cx| {
435                    cell.set_selected(is_selected)
436                        .set_cell_position(cell_position);
437                });
438                cell.clone().into_any_element()
439            }
440            Cell::Raw(cell) => {
441                cell.update(cx, |cell, _cx| {
442                    cell.set_selected(is_selected)
443                        .set_cell_position(cell_position);
444                });
445                cell.clone().into_any_element()
446            }
447        }
448    }
449}
450
451impl Render for NotebookEditor {
452    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
453        div()
454            .key_context("notebook")
455            .track_focus(&self.focus_handle)
456            .on_action(cx.listener(|this, &OpenNotebook, cx| this.open_notebook(&OpenNotebook, cx)))
457            .on_action(cx.listener(|this, &ClearOutputs, cx| this.clear_outputs(cx)))
458            .on_action(cx.listener(|this, &RunAll, cx| this.run_cells(cx)))
459            .on_action(cx.listener(|this, &MoveCellUp, cx| this.move_cell_up(cx)))
460            .on_action(cx.listener(|this, &MoveCellDown, cx| this.move_cell_down(cx)))
461            .on_action(cx.listener(|this, &AddMarkdownBlock, cx| this.add_markdown_block(cx)))
462            .on_action(cx.listener(|this, &AddCodeBlock, cx| this.add_code_block(cx)))
463            .on_action(cx.listener(Self::select_next))
464            .on_action(cx.listener(Self::select_previous))
465            .on_action(cx.listener(Self::select_first))
466            .on_action(cx.listener(Self::select_last))
467            .flex()
468            .items_start()
469            .size_full()
470            .overflow_hidden()
471            .px(Spacing::XLarge.px(cx))
472            .gap(Spacing::XLarge.px(cx))
473            .bg(cx.theme().colors().tab_bar_background)
474            .child(
475                v_flex()
476                    .id("notebook-cells")
477                    .flex_1()
478                    .size_full()
479                    .overflow_y_scroll()
480                    .child(list(self.cell_list.clone()).size_full()),
481            )
482            .child(self.render_notebook_controls(cx))
483    }
484}
485
486impl FocusableView for NotebookEditor {
487    fn focus_handle(&self, _: &AppContext) -> FocusHandle {
488        self.focus_handle.clone()
489    }
490}
491
492// Intended to be a NotebookBuffer
493pub struct NotebookItem {
494    path: PathBuf,
495    project_path: ProjectPath,
496    languages: Arc<LanguageRegistry>,
497    // Raw notebook data
498    notebook: nbformat::v4::Notebook,
499    // Store our version of the notebook in memory (cell_order, cell_map)
500    id: ProjectEntryId,
501}
502
503impl project::Item for NotebookItem {
504    fn try_open(
505        project: &Model<Project>,
506        path: &ProjectPath,
507        cx: &mut AppContext,
508    ) -> Option<Task<gpui::Result<Model<Self>>>> {
509        let path = path.clone();
510        let project = project.clone();
511        let fs = project.read(cx).fs().clone();
512        let languages = project.read(cx).languages().clone();
513
514        if path.path.extension().unwrap_or_default() == "ipynb" {
515            Some(cx.spawn(|mut cx| async move {
516                let abs_path = project
517                    .read_with(&cx, |project, cx| project.absolute_path(&path, cx))?
518                    .ok_or_else(|| anyhow::anyhow!("Failed to find the absolute path"))?;
519
520                // todo: watch for changes to the file
521                let file_content = fs.load(&abs_path.as_path()).await?;
522                let notebook = nbformat::parse_notebook(&file_content);
523
524                let notebook = match notebook {
525                    Ok(nbformat::Notebook::V4(notebook)) => notebook,
526                    // 4.1 - 4.4 are converted to 4.5
527                    Ok(nbformat::Notebook::Legacy(legacy_notebook)) => {
528                        // todo!(): Decide if we want to mutate the notebook by including Cell IDs
529                        // and any other conversions
530                        let notebook = nbformat::upgrade_legacy_notebook(legacy_notebook)?;
531                        notebook
532                    }
533                    // Bad notebooks and notebooks v4.0 and below are not supported
534                    Err(e) => {
535                        anyhow::bail!("Failed to parse notebook: {:?}", e);
536                    }
537                };
538
539                let id = project
540                    .update(&mut cx, |project, cx| project.entry_for_path(&path, cx))?
541                    .context("Entry not found")?
542                    .id;
543
544                cx.new_model(|_| NotebookItem {
545                    path: abs_path,
546                    project_path: path,
547                    languages,
548                    notebook,
549                    id,
550                })
551            }))
552        } else {
553            None
554        }
555    }
556
557    fn entry_id(&self, _: &AppContext) -> Option<ProjectEntryId> {
558        Some(self.id)
559    }
560
561    fn project_path(&self, _: &AppContext) -> Option<ProjectPath> {
562        Some(self.project_path.clone())
563    }
564}
565
566impl NotebookItem {
567    pub fn language_name(&self) -> Option<String> {
568        self.notebook
569            .metadata
570            .language_info
571            .as_ref()
572            .map(|l| l.name.clone())
573            .or(self
574                .notebook
575                .metadata
576                .kernelspec
577                .as_ref()
578                .and_then(|spec| spec.language.clone()))
579    }
580
581    pub fn notebook_language(&self) -> impl Future<Output = Option<Arc<Language>>> {
582        let language_name = self.language_name();
583        let languages = self.languages.clone();
584
585        async move {
586            if let Some(language_name) = language_name {
587                languages.language_for_name(&language_name).await.ok()
588            } else {
589                None
590            }
591        }
592    }
593}
594
595impl EventEmitter<()> for NotebookEditor {}
596
597// pub struct NotebookControls {
598//     pane_focused: bool,
599//     active_item: Option<Box<dyn ItemHandle>>,
600//     // subscription: Option<Subscription>,
601// }
602
603// impl NotebookControls {
604//     pub fn new() -> Self {
605//         Self {
606//             pane_focused: false,
607//             active_item: Default::default(),
608//             // subscription: Default::default(),
609//         }
610//     }
611// }
612
613// impl EventEmitter<ToolbarItemEvent> for NotebookControls {}
614
615// impl Render for NotebookControls {
616//     fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
617//         div().child("notebook controls")
618//     }
619// }
620
621// impl ToolbarItemView for NotebookControls {
622//     fn set_active_pane_item(
623//         &mut self,
624//         active_pane_item: Option<&dyn workspace::ItemHandle>,
625//         cx: &mut ViewContext<Self>,
626//     ) -> workspace::ToolbarItemLocation {
627//         cx.notify();
628//         self.active_item = None;
629
630//         let Some(item) = active_pane_item else {
631//             return ToolbarItemLocation::Hidden;
632//         };
633
634//         ToolbarItemLocation::PrimaryLeft
635//     }
636
637//     fn pane_focus_update(&mut self, pane_focused: bool, _: &mut ViewContext<Self>) {
638//         self.pane_focused = pane_focused;
639//     }
640// }
641
642impl Item for NotebookEditor {
643    type Event = ();
644
645    fn clone_on_split(
646        &self,
647        _workspace_id: Option<workspace::WorkspaceId>,
648        cx: &mut ViewContext<Self>,
649    ) -> Option<gpui::View<Self>>
650    where
651        Self: Sized,
652    {
653        Some(cx.new_view(|cx| Self::new(self.project.clone(), self.notebook_item.clone(), cx)))
654    }
655
656    fn for_each_project_item(
657        &self,
658        cx: &AppContext,
659        f: &mut dyn FnMut(gpui::EntityId, &dyn project::Item),
660    ) {
661        f(self.notebook_item.entity_id(), self.notebook_item.read(cx))
662    }
663
664    fn is_singleton(&self, _cx: &AppContext) -> bool {
665        true
666    }
667
668    fn tab_content(&self, params: TabContentParams, cx: &WindowContext) -> AnyElement {
669        let path = &self.notebook_item.read(cx).path;
670        let title = path
671            .file_name()
672            .unwrap_or_else(|| path.as_os_str())
673            .to_string_lossy()
674            .to_string();
675        Label::new(title)
676            .single_line()
677            .color(params.text_color())
678            .italic(params.preview)
679            .into_any_element()
680    }
681
682    fn tab_icon(&self, _cx: &ui::WindowContext) -> Option<Icon> {
683        Some(IconName::Book.into())
684    }
685
686    fn show_toolbar(&self) -> bool {
687        false
688    }
689
690    // TODO
691    fn pixel_position_of_cursor(&self, _: &AppContext) -> Option<Point<Pixels>> {
692        None
693    }
694
695    // TODO
696    fn as_searchable(&self, _: &View<Self>) -> Option<Box<dyn SearchableItemHandle>> {
697        None
698    }
699
700    fn set_nav_history(&mut self, _: workspace::ItemNavHistory, _: &mut ViewContext<Self>) {
701        // TODO
702    }
703
704    // TODO
705    fn can_save(&self, _cx: &AppContext) -> bool {
706        false
707    }
708    // TODO
709    fn save(
710        &mut self,
711        _format: bool,
712        _project: Model<Project>,
713        _cx: &mut ViewContext<Self>,
714    ) -> Task<Result<()>> {
715        unimplemented!("save() must be implemented if can_save() returns true")
716    }
717
718    // TODO
719    fn save_as(
720        &mut self,
721        _project: Model<Project>,
722        _path: ProjectPath,
723        _cx: &mut ViewContext<Self>,
724    ) -> Task<Result<()>> {
725        unimplemented!("save_as() must be implemented if can_save() returns true")
726    }
727    // TODO
728    fn reload(
729        &mut self,
730        _project: Model<Project>,
731        _cx: &mut ViewContext<Self>,
732    ) -> Task<Result<()>> {
733        unimplemented!("reload() must be implemented if can_save() returns true")
734    }
735
736    fn is_dirty(&self, cx: &AppContext) -> bool {
737        // self.is_dirty(cx) TODO
738        false
739    }
740}
741
742// TODO: Implement this to allow us to persist to the database, etc:
743// impl SerializableItem for NotebookEditor {}
744
745impl ProjectItem for NotebookEditor {
746    type Item = NotebookItem;
747
748    fn for_project_item(
749        project: Model<Project>,
750        item: Model<Self::Item>,
751        cx: &mut ViewContext<Self>,
752    ) -> Self
753    where
754        Self: Sized,
755    {
756        Self::new(project, item, cx)
757    }
758}