1734 lines · 62.7 KB · Rust
RawDownload
#![allow(unused, dead_code)]use std::future::Future;use std::{path::PathBuf, sync::Arc}; use anyhow::{Context as _, Result};use client::proto::ViewId;use collections::HashMap;use feature_flags::{FeatureFlagAppExt as _, NotebookFeatureFlag};use futures::FutureExt;use futures::future::Shared;use gpui::{    AnyElement, App, Entity, EventEmitter, FocusHandle, Focusable, ListScrollEvent, ListState,    Point, Task, actions, list, prelude::*,};use jupyter_protocol::JupyterKernelspec;use language::{Language, LanguageRegistry};use project::{Project, ProjectEntryId, ProjectPath};use settings::Settings as _;use ui::{CommonAnimationExt, Tooltip, prelude::*};use workspace::item::{ItemEvent, SaveOptions, TabContentParams};use workspace::searchable::SearchableItemHandle;use workspace::{Item, ItemHandle, Pane, ProjectItem, ToolbarItemLocation}; use super::{Cell, CellEvent, CellPosition, MarkdownCellEvent, RenderableCell}; use nbformat::v4::CellId;use nbformat::v4::Metadata as NotebookMetadata;use serde_json;use uuid::Uuid; use crate::components::{KernelPickerDelegate, KernelSelector};use crate::kernels::{    Kernel, KernelSession, KernelSpecification, KernelStatus, LocalKernelSpecification,    NativeRunningKernel, RemoteRunningKernel,};use crate::repl_store::ReplStore; use picker::Picker;use runtimelib::{ExecuteRequest, JupyterMessage, JupyterMessageContent};use ui::PopoverMenuHandle;use zed_actions::editor::{MoveDown, MoveUp}; actions!(    notebook,    [        /// Opens a Jupyter notebook file.        OpenNotebook,        /// Runs all cells in the notebook.        RunAll,        /// Runs the current cell.        Run,        /// Clears all cell outputs.        ClearOutputs,        /// Moves the current cell up.        MoveCellUp,        /// Moves the current cell down.        MoveCellDown,        /// Adds a new markdown cell.        AddMarkdownBlock,        /// Adds a new code cell.        AddCodeBlock,        /// Restarts the kernel.        RestartKernel,        /// Interrupts the current execution.        InterruptKernel,    ]); pub(crate) const MAX_TEXT_BLOCK_WIDTH: f32 = 9999.0;pub(crate) const SMALL_SPACING_SIZE: f32 = 8.0;pub(crate) const MEDIUM_SPACING_SIZE: f32 = 12.0;pub(crate) const LARGE_SPACING_SIZE: f32 = 16.0;pub(crate) const GUTTER_WIDTH: f32 = 19.0;pub(crate) const CODE_BLOCK_INSET: f32 = MEDIUM_SPACING_SIZE;pub(crate) const CONTROL_SIZE: f32 = 20.0; pub fn init(cx: &mut App) {    if cx.has_flag::<NotebookFeatureFlag>() || std::env::var("LOCAL_NOTEBOOK_DEV").is_ok() {        workspace::register_project_item::<NotebookEditor>(cx);    }     cx.observe_flag::<NotebookFeatureFlag, _>({        move |is_enabled, cx| {            if is_enabled {                workspace::register_project_item::<NotebookEditor>(cx);            } else {                // todo: there is no way to unregister a project item, so if the feature flag                // gets turned off they need to restart Zed.            }        }    })    .detach();} pub struct NotebookEditor {    languages: Arc<LanguageRegistry>,    project: Entity<Project>,    worktree_id: project::WorktreeId,     focus_handle: FocusHandle,    notebook_item: Entity<NotebookItem>,    notebook_language: Shared<Task<Option<Arc<Language>>>>,     remote_id: Option<ViewId>,    cell_list: ListState,     selected_cell_index: usize,    cell_order: Vec<CellId>,    original_cell_order: Vec<CellId>,    cell_map: HashMap<CellId, Cell>,    kernel: Kernel,    kernel_specification: Option<KernelSpecification>,    execution_requests: HashMap<String, CellId>,    kernel_picker_handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>,} impl NotebookEditor {    pub fn new(        project: Entity<Project>,        notebook_item: Entity<NotebookItem>,        window: &mut Window,        cx: &mut Context<Self>,    ) -> Self {        let focus_handle = cx.focus_handle();         let languages = project.read(cx).languages().clone();        let language_name = notebook_item.read(cx).language_name();        let worktree_id = notebook_item.read(cx).project_path.worktree_id;         let notebook_language = notebook_item.read(cx).notebook_language();        let notebook_language = cx            .spawn_in(window, async move |_, _| notebook_language.await)            .shared();         let mut cell_order = vec![]; // Vec<CellId>        let mut cell_map = HashMap::default(); // HashMap<CellId, Cell>         let cell_count = notebook_item.read(cx).notebook.cells.len();        for index in 0..cell_count {            let cell = notebook_item.read(cx).notebook.cells[index].clone();            let cell_id = cell.id();            cell_order.push(cell_id.clone());            let cell_entity = Cell::load(&cell, &languages, notebook_language.clone(), window, cx);             match &cell_entity {                Cell::Code(code_cell) => {                    let cell_id_for_focus = cell_id.clone();                    cx.subscribe(code_cell, move |this, cell, event, cx| match event {                        CellEvent::Run(cell_id) => this.execute_cell(cell_id.clone(), cx),                        CellEvent::FocusedIn(_) => {                            if let Some(index) = this                                .cell_order                                .iter()                                .position(|id| id == &cell_id_for_focus)                            {                                this.selected_cell_index = index;                                cx.notify();                            }                        }                    })                    .detach();                     let cell_id_for_editor = cell_id.clone();                    let editor = code_cell.read(cx).editor().clone();                    cx.subscribe(&editor, move |this, _editor, event, cx| {                        if let editor::EditorEvent::Focused = event {                            if let Some(index) = this                                .cell_order                                .iter()                                .position(|id| id == &cell_id_for_editor)                            {                                this.selected_cell_index = index;                                cx.notify();                            }                        }                    })                    .detach();                }                Cell::Markdown(markdown_cell) => {                    let cell_id_for_focus = cell_id.clone();                    cx.subscribe(                        markdown_cell,                        move |_this, cell, event: &MarkdownCellEvent, cx| {                            match event {                                MarkdownCellEvent::FinishedEditing => {                                    cell.update(cx, |cell, cx| {                                        cell.reparse_markdown(cx);                                    });                                }                                MarkdownCellEvent::Run(_cell_id) => {                                    // run is handled separately by move_to_next_cell                                    // Just reparse here                                    cell.update(cx, |cell, cx| {                                        cell.reparse_markdown(cx);                                    });                                }                            }                        },                    )                    .detach();                     let cell_id_for_editor = cell_id.clone();                    let editor = markdown_cell.read(cx).editor().clone();                    cx.subscribe(&editor, move |this, _editor, event, cx| {                        if let editor::EditorEvent::Focused = event {                            if let Some(index) = this                                .cell_order                                .iter()                                .position(|id| id == &cell_id_for_editor)                            {                                this.selected_cell_index = index;                                cx.notify();                            }                        }                    })                    .detach();                }                Cell::Raw(_) => {}            }             cell_map.insert(cell_id.clone(), cell_entity);        }         let notebook_handle = cx.entity().downgrade();        let cell_count = cell_order.len();         let this = cx.entity();        let cell_list = ListState::new(cell_count, gpui::ListAlignment::Top, px(1000.));         let mut editor = Self {            project,            languages: languages.clone(),            worktree_id,            focus_handle,            notebook_item: notebook_item.clone(),            notebook_language,            remote_id: None,            cell_list,            selected_cell_index: 0,            cell_order: cell_order.clone(),            original_cell_order: cell_order.clone(),            cell_map: cell_map.clone(),            kernel: Kernel::Shutdown, // TODO: use recommended kernel after the implementation is done in repl            kernel_specification: None,            execution_requests: HashMap::default(),            kernel_picker_handle: PopoverMenuHandle::default(),        };        editor.launch_kernel(window, cx);        editor.refresh_language(cx);         cx.subscribe(&notebook_item, |this, _item, _event, cx| {            this.refresh_language(cx);        })        .detach();         editor    }     fn refresh_language(&mut self, cx: &mut Context<Self>) {        let notebook_language = self.notebook_item.read(cx).notebook_language();        let task = cx.spawn(async move |this, cx| {            let language = notebook_language.await;            if let Some(this) = this.upgrade() {                this.update(cx, |this, cx| {                    for cell in this.cell_map.values() {                        if let Cell::Code(code_cell) = cell {                            code_cell.update(cx, |cell, cx| {                                cell.set_language(language.clone(), cx);                            });                        }                    }                });            }            language        });        self.notebook_language = task.shared();    }     fn has_structural_changes(&self) -> bool {        self.cell_order != self.original_cell_order    }     fn has_content_changes(&self, cx: &App) -> bool {        self.cell_map.values().any(|cell| cell.is_dirty(cx))    }     pub fn to_notebook(&self, cx: &App) -> nbformat::v4::Notebook {        let cells: Vec<nbformat::v4::Cell> = self            .cell_order            .iter()            .filter_map(|cell_id| {                self.cell_map                    .get(cell_id)                    .map(|cell| cell.to_nbformat_cell(cx))            })            .collect();         let metadata = self.notebook_item.read(cx).notebook.metadata.clone();         nbformat::v4::Notebook {            metadata,            nbformat: 4,            nbformat_minor: 5,            cells,        }    }     pub fn mark_as_saved(&mut self, cx: &mut Context<Self>) {        self.original_cell_order = self.cell_order.clone();         for cell in self.cell_map.values() {            match cell {                Cell::Code(code_cell) => {                    code_cell.update(cx, |code_cell, cx| {                        let editor = code_cell.editor();                        editor.update(cx, |editor, cx| {                            editor.buffer().update(cx, |buffer, cx| {                                if let Some(buf) = buffer.as_singleton() {                                    buf.update(cx, |b, cx| {                                        let version = b.version();                                        b.did_save(version, None, cx);                                    });                                }                            });                        });                    });                }                Cell::Markdown(markdown_cell) => {                    markdown_cell.update(cx, |markdown_cell, cx| {                        let editor = markdown_cell.editor();                        editor.update(cx, |editor, cx| {                            editor.buffer().update(cx, |buffer, cx| {                                if let Some(buf) = buffer.as_singleton() {                                    buf.update(cx, |b, cx| {                                        let version = b.version();                                        b.did_save(version, None, cx);                                    });                                }                            });                        });                    });                }                Cell::Raw(_) => {}            }        }        cx.notify();    }     fn launch_kernel(&mut self, window: &mut Window, cx: &mut Context<Self>) {        // use default Python kernel if no specification is set        let spec = self.kernel_specification.clone().unwrap_or_else(|| {            KernelSpecification::Jupyter(LocalKernelSpecification {                name: "python3".to_string(),                path: PathBuf::from("python3"),                kernelspec: JupyterKernelspec {                    argv: vec![                        "python3".to_string(),                        "-m".to_string(),                        "ipykernel_launcher".to_string(),                        "-f".to_string(),                        "{connection_file}".to_string(),                    ],                    display_name: "Python 3".to_string(),                    language: "python".to_string(),                    interrupt_mode: None,                    metadata: None,                    env: None,                },            })        });         self.launch_kernel_with_spec(spec, window, cx);    }     fn launch_kernel_with_spec(        &mut self,        spec: KernelSpecification,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let entity_id = cx.entity_id();        let working_directory = self            .project            .read(cx)            .worktrees(cx)            .next()            .map(|worktree| worktree.read(cx).abs_path().to_path_buf())            .unwrap_or_else(std::env::temp_dir);        let fs = self.project.read(cx).fs().clone();        let view = cx.entity();         self.kernel_specification = Some(spec.clone());         self.notebook_item.update(cx, |item, cx| {            let kernel_name = spec.name().to_string();            let language = spec.language().to_string();             let display_name = match &spec {                KernelSpecification::Jupyter(s) => s.kernelspec.display_name.clone(),                KernelSpecification::PythonEnv(s) => s.kernelspec.display_name.clone(),                KernelSpecification::Remote(s) => s.kernelspec.display_name.clone(),            };             let kernelspec_json = serde_json::json!({                "display_name": display_name,                "name": kernel_name,                "language": language            });             if let Ok(k) = serde_json::from_value(kernelspec_json) {                item.notebook.metadata.kernelspec = Some(k);                cx.emit(());            }        });         let kernel_task = match spec {            KernelSpecification::Jupyter(local_spec) => NativeRunningKernel::new(                local_spec,                entity_id,                working_directory,                fs,                view,                window,                cx,            ),            KernelSpecification::PythonEnv(env_spec) => NativeRunningKernel::new(                env_spec.as_local_spec(),                entity_id,                working_directory,                fs,                view,                window,                cx,            ),            KernelSpecification::Remote(remote_spec) => {                RemoteRunningKernel::new(remote_spec, working_directory, view, window, cx)            }        };         let pending_kernel = cx            .spawn(async move |this, cx| {                let kernel = kernel_task.await;                 match kernel {                    Ok(kernel) => {                        this.update(cx, |editor, cx| {                            editor.kernel = Kernel::RunningKernel(kernel);                            cx.notify();                        })                        .ok();                    }                    Err(err) => {                        this.update(cx, |editor, cx| {                            editor.kernel = Kernel::ErroredLaunch(err.to_string());                            cx.notify();                        })                        .ok();                    }                }            })            .shared();         self.kernel = Kernel::StartingKernel(pending_kernel);        cx.notify();    }     // Note: Python environments are only detected as kernels if ipykernel is installed.    // Users need to run `pip install ipykernel` (or `uv pip install ipykernel`) in their    // virtual environment for it to appear in the kernel selector.    // This happens because we have an ipykernel check inside the function python_env_kernel_specification in mod.rs L:121     fn change_kernel(        &mut self,        spec: KernelSpecification,        window: &mut Window,        cx: &mut Context<Self>,    ) {        if let Kernel::RunningKernel(kernel) = &mut self.kernel {            kernel.force_shutdown(window, cx).detach();        }         self.execution_requests.clear();         self.launch_kernel_with_spec(spec, window, cx);    }     fn restart_kernel(&mut self, _: &RestartKernel, window: &mut Window, cx: &mut Context<Self>) {        if let Some(spec) = self.kernel_specification.clone() {            if let Kernel::RunningKernel(kernel) = &mut self.kernel {                kernel.force_shutdown(window, cx).detach();            }             self.kernel = Kernel::Restarting;            cx.notify();             self.launch_kernel_with_spec(spec, window, cx);        }    }     fn interrupt_kernel(

Lines 1–500 of 1734.