Address additional suggestions

crates/ruff_server/src/fix.rs

@@ -70,26 +70,22 @@ pub(crate) fn fix_all(
     if let (Some(source_notebook), Some(modified_notebook)) =
         (source_kind.as_ipy_notebook(), transformed.as_ipy_notebook())
     {
+        fn cell_source(cell: &ruff_notebook::Cell) -> String {
+            match cell.source() {
+                SourceValue::String(string) => string.clone(),
+                SourceValue::StringArray(array) => array.join(""),
+            }
+        }
+
         let Some(notebook) = query.as_notebook() else {
             anyhow::bail!("Notebook document expected from notebook source kind");
         };
         let mut fixes = Fixes::default();
         for ((source, modified), url) in source_notebook
             .cells()
             .iter()
-            .map(|cell| match cell.source() {
-                SourceValue::String(string) => string.clone(),
-                SourceValue::StringArray(array) => array.join(""),
-            })
-            .zip(
-                modified_notebook
-                    .cells()
-                    .iter()
-                    .map(|cell| match cell.source() {
-                        SourceValue::String(string) => string.clone(),
-                        SourceValue::StringArray(array) => array.join(""),
-                    }),
-            )
+            .map(cell_source)
+            .zip(modified_notebook.cells().iter().map(cell_source))
             .zip(notebook.urls())
         {
             let source_index = LineIndex::from_source_text(&source);

crates/ruff_server/src/lint.rs

@@ -16,7 +16,7 @@ use ruff_python_index::Indexer;
 use ruff_python_parser::lexer::LexResult;
 use ruff_python_parser::AsMode;
 use ruff_source_file::{LineIndex, Locator};
-use ruff_text_size::Ranged;
+use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
@@ -129,20 +129,23 @@ pub(crate) fn check(
         }
     }
 
-    for (index, diagnostic) in data
+    let lsp_diagnostics = data
         .into_iter()
         .zip(noqa_edits)
         .map(|(diagnostic, noqa_edit)| {
             to_lsp_diagnostic(diagnostic, &noqa_edit, &source_kind, &index, encoding)
-        })
-    {
-        if let Some(notebook) = query.as_notebook() {
+        });
+
+    if let Some(notebook) = query.as_notebook() {
+        for (index, diagnostic) in lsp_diagnostics {
             let Some(uri) = notebook.cell_uri_by_index(index) else {
                 tracing::warn!("Unable to find notebook cell at index {index}.");
                 continue;
             };
             diagnostics.entry(uri.clone()).or_default().push(diagnostic);
-        } else {
+        }
+    } else {
+        for (_, diagnostic) in lsp_diagnostics {
             diagnostics
                 .entry(query.make_key().into_url())
                 .or_default()
@@ -210,42 +213,12 @@ fn to_lsp_diagnostic(
                 .into_iter()
                 .flat_map(Fix::edits)
                 .map(|edit| lsp_types::TextEdit {
-                    range: if let Some(notebook_index) =
-                        source_kind.as_ipy_notebook().map(Notebook::index)
-                    {
-                        edit.range()
-                            .to_notebook_range(
-                                source_kind.source_code(),
-                                index,
-                                notebook_index,
-                                encoding,
-                            )
-                            .range
-                    } else {
-                        edit.range()
-                            .to_range(source_kind.source_code(), index, encoding)
-                    },
+                    range: diagnostic_edit_range(edit.range(), source_kind, index, encoding),
                     new_text: edit.content().unwrap_or_default().to_string(),
                 })
                 .collect();
             let noqa_edit = noqa_edit.as_ref().map(|noqa_edit| lsp_types::TextEdit {
-                range: if let Some(notebook_index) =
-                    source_kind.as_ipy_notebook().map(Notebook::index)
-                {
-                    noqa_edit
-                        .range()
-                        .to_notebook_range(
-                            source_kind.source_code(),
-                            index,
-                            notebook_index,
-                            encoding,
-                        )
-                        .range
-                } else {
-                    noqa_edit
-                        .range()
-                        .to_range(source_kind.source_code(), index, encoding)
-                },
+                range: diagnostic_edit_range(noqa_edit.range(), source_kind, index, encoding),
                 new_text: noqa_edit.content().unwrap_or_default().to_string(),
             });
             serde_json::to_value(AssociatedDiagnosticData {
@@ -295,6 +268,21 @@ fn to_lsp_diagnostic(
     )
 }
 
+fn diagnostic_edit_range(
+    range: TextRange,
+    source_kind: &SourceKind,
+    index: &LineIndex,
+    encoding: PositionEncoding,
+) -> lsp_types::Range {
+    if let Some(notebook_index) = source_kind.as_ipy_notebook().map(Notebook::index) {
+        range
+            .to_notebook_range(source_kind.source_code(), index, notebook_index, encoding)
+            .range
+    } else {
+        range.to_range(source_kind.source_code(), index, encoding)
+    }
+}
+
 fn severity(code: &str) -> lsp_types::DiagnosticSeverity {
     match code {
         // F821: undefined name <name>

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

@@ -3,6 +3,7 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
+use lsp_server::ErrorCode;
 use lsp_types as types;
 use lsp_types::notification as notif;
 
@@ -26,11 +27,13 @@ impl super::SyncNotificationHandler for DidChange {
             content_changes,
         }: types::DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let key = session.key_from_url(&uri);
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
 
         session
             .update_text_document(&key, content_changes, new_version)
-            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+            .with_failure_code(ErrorCode::InternalError)?;
 
         // Publish diagnostics if the client doesnt support pull diagnostics
         if !session.resolved_client_capabilities().pull_diagnostics {

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

@@ -23,7 +23,9 @@ impl super::SyncNotificationHandler for DidChangeNotebook {
             change: types::NotebookDocumentChangeEvent { cells, metadata },
         }: types::DidChangeNotebookDocumentParams,
     ) -> Result<()> {
-        let key = session.key_from_url(&uri);
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
         session
             .update_notebook_document(&key, cells, metadata, version)
             .with_failure_code(ErrorCode::InternalError)?;

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

@@ -22,7 +22,9 @@ impl super::SyncNotificationHandler for DidCloseNotebook {
             ..
         }: types::DidCloseNotebookDocumentParams,
     ) -> Result<()> {
-        let key = session.key_from_url(&uri);
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
 
         session
             .close_document(&key)

crates/ruff_server/src/session.rs

@@ -7,6 +7,7 @@ mod settings;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use anyhow::anyhow;
 use lsp_types::{ClientCapabilities, NotebookDocumentCellChange, Url};
 
 use crate::edit::{DocumentKey, DocumentVersion, NotebookDocument};
@@ -54,13 +55,15 @@ impl Session {
         }
     }
 
-    pub(crate) fn key_from_url(&self, url: &lsp_types::Url) -> DocumentKey {
-        self.index.key_from_url(url)
+    pub(crate) fn key_from_url(&self, url: &lsp_types::Url) -> crate::Result<DocumentKey> {
+        self.index
+            .key_from_url(url)
+            .ok_or_else(|| anyhow!("No document found for {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 = self.key_from_url(url);
+        let key = self.key_from_url(url).ok()?;
         Some(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
             client_settings: self.index.client_settings(&key, &self.global_settings),

crates/ruff_server/src/session/index.rs

@@ -48,7 +48,7 @@ enum DocumentController {
 }
 
 /// A read-only query to an open document.
-/// This query can 'select' a text document, full notebook, or notebook cell.
+/// This query can 'select' a text document, full notebook, or a specific notebook cell.
 /// It also includes document settings.
 #[derive(Clone)]
 pub(crate) enum DocumentQuery {
@@ -58,8 +58,7 @@ pub(crate) enum DocumentQuery {
         settings: Arc<RuffSettings>,
     },
     Notebook {
-        /// The selected notebook cell.
-        /// If this is `None`, the entire notebook file is selected.
+        /// The selected notebook cell, if it exists.
         cell_uri: Option<lsp_types::Url>,
         file_path: PathBuf,
         notebook: Arc<NotebookDocument>,
@@ -96,12 +95,7 @@ impl Index {
         new_version: DocumentVersion,
         encoding: PositionEncoding,
     ) -> crate::Result<()> {
-        let Some(path) = self.path_for_key(key).cloned() else {
-            anyhow::bail!("Tried to open unavailable document `{key}`");
-        };
-        let Some(controller) = self.documents.get_mut(&path) else {
-            anyhow::bail!("Document controller not available at `{}`", path.display());
-        };
+        let controller = self.document_controller_for_key(key)?;
         let Some(document) = controller.as_text_mut() else {
             anyhow::bail!("Text document URI does not point to a text document");
         };
@@ -116,22 +110,22 @@ impl Index {
         Ok(())
     }
 
-    pub(super) fn key_from_url(&self, url: &lsp_types::Url) -> DocumentKey {
+    pub(super) fn key_from_url(&self, url: &lsp_types::Url) -> Option<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),
+            return Some(DocumentKey::NotebookCell(url.clone()));
         }
+        let path = url.to_file_path().ok()?;
+        Some(
+            match path
+                .extension()
+                .unwrap_or_default()
+                .to_str()
+                .unwrap_or_default()
+            {
+                "ipynb" => DocumentKey::Notebook(path),
+                _ => DocumentKey::Text(path),
+            },
+        )
     }
 
     pub(super) fn update_notebook_document(
@@ -142,17 +136,17 @@ impl Index {
         new_version: DocumentVersion,
         encoding: PositionEncoding,
     ) -> crate::Result<()> {
-        let Some(path) = self.path_for_key(key).cloned() else {
-            anyhow::bail!("Tried to open unavailable document `{key}`");
-        };
-
         // update notebook cell index
         if let Some(lsp_types::NotebookDocumentCellChangeStructure {
             did_open,
             did_close,
             ..
         }) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
         {
+            let Some(path) = self.path_for_key(key).cloned() else {
+                anyhow::bail!("Tried to open unavailable document `{key}`");
+            };
+
             for opened_cell in did_open.iter().flatten() {
                 self.notebook_cells
                     .insert(opened_cell.uri.clone(), path.clone());
@@ -167,9 +161,7 @@ impl Index {
             }
         }
 
-        let Some(controller) = self.documents.get_mut(&path) else {
-            anyhow::bail!("Document controller not available at `{}`", path.display());
-        };
+        let controller = self.document_controller_for_key(key)?;
         let Some(notebook) = controller.as_notebook_mut() else {
             anyhow::bail!("Notebook document URI does not point to a notebook document");
         };
@@ -306,6 +298,19 @@ impl Index {
         client_settings.clone()
     }
 
+    fn document_controller_for_key(
+        &mut self,
+        key: &DocumentKey,
+    ) -> crate::Result<&mut DocumentController> {
+        let Some(path) = self.path_for_key(key).cloned() else {
+            anyhow::bail!("Tried to open unavailable document `{key}`");
+        };
+        let Some(controller) = self.documents.get_mut(&path) else {
+            anyhow::bail!("Document controller not available at `{}`", path.display());
+        };
+        Ok(controller)
+    }
+
     fn path_for_key<'a>(&'a self, key: &'a DocumentKey) -> Option<&'a PathBuf> {
         match key {
             DocumentKey::Notebook(path) | DocumentKey::Text(path) => Some(path),