Address some suggestions

Author: snowsignal

crates/ruff_notebook/src/notebook.rs

@@ -85,23 +85,6 @@ impl Notebook {
         Self::from_reader(Cursor::new(source_code))
     }
 
-    /// Generate a pseudo-representation of a notebook that can be used
-    /// for linting by the language server. As this is not generated directly from the raw JSON
-    /// of a notebook file, writing this back into the file system is a bad idea.
-    pub fn from_cells(
-        cells: Vec<Cell>,
-        metadata: crate::RawNotebookMetadata,
-    ) -> Result<Self, NotebookError> {
-        let raw_notebook = RawNotebook {
-            cells,
-            metadata,
-            nbformat: 4,
-            nbformat_minor: 5,
-        };
-
-        Self::from_raw(raw_notebook, false)
-    }
-
     /// Read a Jupyter Notebook from a [`Read`] implementer.
     ///
     /// See also the black implementation
@@ -130,10 +113,10 @@ impl Notebook {
                 });
             }
         };
-        Self::from_raw(raw_notebook, trailing_newline)
+        Self::from_raw_notebook(raw_notebook, trailing_newline)
     }
 
-    fn from_raw(
+    pub fn from_raw_notebook(
         mut raw_notebook: RawNotebook,
         trailing_newline: bool,
     ) -> Result<Self, NotebookError> {

crates/ruff_server/src/edit.rs

@@ -5,13 +5,13 @@ mod notebook;
 mod range;
 mod replacement;
 
-use std::{collections::HashMap, ffi::OsStr, path::PathBuf};
+use std::{collections::HashMap, path::PathBuf};
 
 pub(crate) use document::DocumentVersion;
 pub use document::TextDocument;
 use lsp_types::PositionEncodingKind;
 pub(crate) use notebook::NotebookDocument;
-pub(crate) use range::{RangeExt, ToRangeExt};
+pub(crate) use range::{NotebookRange, RangeExt, ToRangeExt};
 pub(crate) use replacement::Replacement;
 
 use crate::{fix::Fixes, session::ResolvedClientCapabilities};
@@ -41,26 +41,6 @@ pub(crate) enum DocumentKey {
 }
 
 impl DocumentKey {
-    /// Creates a document key from a URL provided in an LSP request.
-    pub(crate) fn from_url(url: &lsp_types::Url) -> Self {
-        if url.scheme() != "file" {
-            return Self::NotebookCell(url.clone());
-        }
-        let Some(path) = url.to_file_path().ok() else {
-            return Self::NotebookCell(url.clone());
-        };
-
-        // figure out whether this is a notebook or a text document
-        if path.extension() == Some(OsStr::new("ipynb")) {
-            Self::Notebook(path)
-        } else {
-            // Until we support additional document types, we need to confirm
-            // that any non-notebook file is a Python file
-            debug_assert_eq!(path.extension(), Some(OsStr::new("py")));
-            Self::Text(path)
-        }
-    }
-
     /// Converts the key back into its original URL.
     pub(crate) fn into_url(self) -> lsp_types::Url {
         match self {

crates/ruff_server/src/edit/notebook.rs

@@ -8,6 +8,8 @@ use crate::{PositionEncoding, TextDocument};
 
 use super::DocumentVersion;
 
+pub(super) type CellId = usize;
+
 /// The state of a notebook document in the server. Contains an array of cells whose
 /// contents are internally represented by [`TextDocument`]s.
 #[derive(Clone, Debug)]
@@ -16,9 +18,10 @@ pub(crate) struct NotebookDocument {
     metadata: ruff_notebook::RawNotebookMetadata,
     version: DocumentVersion,
     // Used to quickly find the index of a cell for a given URL.
-    cell_index: FxHashMap<lsp_types::Url, usize>,
+    cell_index: FxHashMap<lsp_types::Url, CellId>,
 }
 
+/// A single cell within a notebook, which has text contents represented as a `TextDocument`.
 #[derive(Clone, Debug)]
 struct NotebookCell {
     url: Url,
@@ -82,9 +85,15 @@ impl NotebookDocument {
                 }
             })
             .collect();
+        let raw_notebook = ruff_notebook::RawNotebook {
+            cells,
+            metadata: self.metadata.clone(),
+            nbformat: 4,
+            nbformat_minor: 5,
+        };
 
-        ruff_notebook::Notebook::from_cells(cells, self.metadata.clone())
-            .expect("notebook should convert successfully")
+        ruff_notebook::Notebook::from_raw_notebook(raw_notebook, false)
+            .unwrap_or_else(|err| panic!("Server notebook document could not be converted to Ruff's notebook document format: {err}"))
     }
 
     pub(crate) fn update(
@@ -103,18 +112,22 @@ impl NotebookDocument {
         }) = cells
         {
             if let Some(structure) = structure {
-                let start = usize::try_from(structure.array.start).unwrap();
-                let delete = usize::try_from(structure.array.delete_count).unwrap();
+                let start = structure.array.start as usize;
+                let delete = structure.array.delete_count as usize;
                 if delete > 0 {
-                    self.cells.drain(start..start + delete);
+                    for cell in self.cells.drain(start..start + delete) {
+                        self.cell_index.remove(&cell.url);
+                    }
                 }
                 for cell in structure.array.cells.into_iter().flatten().rev() {
                     self.cells
                         .insert(start, NotebookCell::new(cell, String::new(), version));
                 }
 
-                // the array has been updated - rebuild the cell index
-                self.rebuild_cell_index();
+                // register any new cells in the index and update existing ones that came after the insertion
+                for (i, cell) in self.cells.iter().enumerate().skip(start) {
+                    self.cell_index.insert(cell.url.clone(), i);
+                }
             }
             if let Some(cell_data) = data {
                 for cell in cell_data {
@@ -138,20 +151,24 @@ impl NotebookDocument {
         Ok(())
     }
 
+    /// Get the current version of the notebook document.
     pub(crate) fn version(&self) -> DocumentVersion {
         self.version
     }
 
-    pub(crate) fn cell_uri_by_index(&self, index: usize) -> Option<&lsp_types::Url> {
+    /// Get the URI for a cell by its index within the cell array.
+    pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
         self.cells.get(index).map(|cell| &cell.url)
     }
 
+    /// Get the text document representing the contents of a cell by the cell URI.
     pub(crate) fn cell_document_by_uri(&self, uri: &lsp_types::Url) -> Option<&TextDocument> {
         self.cells
             .get(*self.cell_index.get(uri)?)
             .map(|cell| &cell.document)
     }
 
+    /// Returns a list of cell URIs in the order they appear in the array.
     pub(crate) fn urls(&self) -> impl Iterator<Item = &lsp_types::Url> {
         self.cells.iter().map(|cell| &cell.url)
     }
@@ -160,11 +177,7 @@ impl NotebookDocument {
         self.cells.get_mut(*self.cell_index.get(uri)?)
     }
 
-    fn rebuild_cell_index(&mut self) {
-        self.cell_index = Self::make_cell_index(&self.cells);
-    }
-
-    fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<lsp_types::Url, usize> {
+    fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<lsp_types::Url, CellId> {
         let mut index =
             HashMap::with_capacity_and_hasher(cells.len(), BuildHasherDefault::default());
         for (i, cell) in cells.iter().enumerate() {

crates/ruff_server/src/edit/range.rs

@@ -1,10 +1,16 @@
+use super::notebook;
 use super::PositionEncoding;
 use lsp_types as types;
 use ruff_notebook::NotebookIndex;
 use ruff_source_file::OneIndexed;
 use ruff_source_file::{LineIndex, SourceLocation};
 use ruff_text_size::{TextRange, TextSize};
 
+pub(crate) struct NotebookRange {
+    pub(crate) cell: notebook::CellId,
+    pub(crate) range: types::Range,
+}
+
 pub(crate) trait RangeExt {
     fn to_text_range(&self, text: &str, index: &LineIndex, encoding: PositionEncoding)
         -> TextRange;
@@ -18,7 +24,7 @@ pub(crate) trait ToRangeExt {
         source_index: &LineIndex,
         notebook_index: &NotebookIndex,
         encoding: PositionEncoding,
-    ) -> (usize, types::Range);
+    ) -> NotebookRange;
 }
 
 fn u32_index_to_usize(index: u32) -> usize {
@@ -112,7 +118,7 @@ impl ToRangeExt for TextRange {
         source_index: &LineIndex,
         notebook_index: &NotebookIndex,
         encoding: PositionEncoding,
-    ) -> (usize, types::Range) {
+    ) -> NotebookRange {
         let start = offset_to_source_location(self.start(), text, source_index, encoding);
         let mut end = offset_to_source_location(self.end(), text, source_index, encoding);
         let cell = notebook_index.cell(start.row);
@@ -131,10 +137,10 @@ impl ToRangeExt for TextRange {
         let start = source_location_to_position(&notebook_index.translate_location(&start));
         let end = source_location_to_position(&notebook_index.translate_location(&end));
 
-        (
-            cell.map(OneIndexed::to_zero_indexed).unwrap_or_default(),
-            types::Range { start, end },
-        )
+        NotebookRange {
+            cell: cell.map(OneIndexed::to_zero_indexed).unwrap_or_default(),
+            range: types::Range { start, end },
+        }
     }
 }
 

crates/ruff_server/src/lint.rs

@@ -19,7 +19,11 @@ use ruff_text_size::Ranged;
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
-use crate::{edit::ToRangeExt, session::DocumentQuery, PositionEncoding, DIAGNOSTIC_NAME};
+use crate::{
+    edit::{NotebookRange, ToRangeExt},
+    session::DocumentQuery,
+    PositionEncoding, DIAGNOSTIC_NAME,
+};
 
 /// This is serialized on the diagnostic `data` field.
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -214,7 +218,7 @@ fn to_lsp_diagnostic(
                                 notebook_index,
                                 encoding,
                             )
-                            .1
+                            .range
                     } else {
                         edit.range()
                             .to_range(source_kind.source_code(), index, encoding)
@@ -234,7 +238,7 @@ fn to_lsp_diagnostic(
                             notebook_index,
                             encoding,
                         )
-                        .1
+                        .range
                 } else {
                     noqa_edit
                         .range()
@@ -258,7 +262,7 @@ fn to_lsp_diagnostic(
     let cell: usize;
 
     if let Some(notebook_index) = source_kind.as_ipy_notebook().map(Notebook::index) {
-        (cell, range) = diagnostic_range.to_notebook_range(
+        NotebookRange { cell, range } = diagnostic_range.to_notebook_range(
             source_kind.source_code(),
             index,
             notebook_index,

crates/ruff_server/src/server/api/notifications/did_change.rs

@@ -1,4 +1,3 @@
-use crate::edit::DocumentKey;
 use crate::server::api::diagnostics::publish_diagnostics_for_document;
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
@@ -27,7 +26,7 @@ impl super::SyncNotificationHandler for DidChange {
             content_changes,
         }: types::DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let key = DocumentKey::from_url(&uri);
+        let key = session.key_from_url(&uri);
 
         session
             .update_text_document(&key, content_changes, new_version)

crates/ruff_server/src/server/api/notifications/did_change_notebook.rs

@@ -1,4 +1,3 @@
-use crate::edit::DocumentKey;
 use crate::server::api::diagnostics::publish_diagnostics_for_document;
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
@@ -24,7 +23,7 @@ impl super::SyncNotificationHandler for DidChangeNotebook {
             change: types::NotebookDocumentChangeEvent { cells, metadata },
         }: types::DidChangeNotebookDocumentParams,
     ) -> Result<()> {
-        let key = DocumentKey::from_url(&uri);
+        let key = session.key_from_url(&uri);
         session
             .update_notebook_document(&key, cells, metadata, version)
             .with_failure_code(ErrorCode::InternalError)?;

crates/ruff_server/src/server/api/notifications/did_close_notebook.rs

@@ -1,4 +1,3 @@
-use crate::edit::DocumentKey;
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
@@ -23,7 +22,7 @@ impl super::SyncNotificationHandler for DidCloseNotebook {
             ..
         }: types::DidCloseNotebookDocumentParams,
     ) -> Result<()> {
-        let key = DocumentKey::from_url(&uri);
+        let key = session.key_from_url(&uri);
 
         session
             .close_document(&key)

crates/ruff_server/src/session.rs

@@ -54,9 +54,13 @@ impl Session {
         }
     }
 
+    pub(crate) fn key_from_url(&self, url: &lsp_types::Url) -> DocumentKey {
+        self.index.key_from_url(url)
+    }
+
     /// Creates a document snapshot with the URL referencing the document to snapshot.
     pub(crate) fn take_snapshot(&self, url: &Url) -> Option<DocumentSnapshot> {
-        let key = DocumentKey::from_url(url);
+        let key = self.key_from_url(url);
         Some(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
             client_settings: self.index.client_settings(&key, &self.global_settings),
@@ -67,7 +71,7 @@ impl Session {
 
     /// Updates a text document at the associated `key`.
     ///
-    /// The document key must point to a text document.
+    /// The document key must point to a text document, or this will throw an error.
     pub(crate) fn update_text_document(
         &mut self,
         key: &DocumentKey,
@@ -83,7 +87,8 @@ impl Session {
     /// Updates a notebook document at the associated `key` with potentially new
     /// cell, metadata, and version values.
     ///
-    /// The document key must point to a notebook document or cell.
+    /// The document key must point to a notebook document or cell, or this will
+    /// throw an error.
     pub(crate) fn update_notebook_document(
         &mut self,
         key: &DocumentKey,

crates/ruff_server/src/session/index.rs

@@ -116,6 +116,24 @@ impl Index {
         Ok(())
     }
 
+    pub(super) fn key_from_url(&self, url: &lsp_types::Url) -> DocumentKey {
+        if self.notebook_cells.contains_key(url) {
+            return DocumentKey::NotebookCell(url.clone());
+        }
+        let path = url
+            .to_file_path()
+            .expect("non-notebook cell URL should be a valid path");
+        match path
+            .extension()
+            .unwrap_or_default()
+            .to_str()
+            .unwrap_or_default()
+        {
+            "ipynb" => DocumentKey::Notebook(path),
+            _ => DocumentKey::Text(path),
+        }
+    }
+
     pub(super) fn update_notebook_document(
         &mut self,
         key: &DocumentKey,