ruff.applyFormat now formats an entire notebook document

crates/ruff_server/src/server/api/requests/execute_command.rs

@@ -73,12 +73,10 @@ impl super::SyncRequestHandler for ExecuteCommand {
                         .with_failure_code(ErrorCode::InternalError)?;
                 }
                 Command::Format => {
-                    let response = super::format::format_document(&snapshot)?;
-                    if let Some(edits) = response {
-                        edit_tracker
-                            .set_edits_for_document(uri, version, edits)
-                            .with_failure_code(ErrorCode::InternalError)?;
-                    }
+                    let fixes = super::format::format_full_document(&snapshot)?;
+                    edit_tracker
+                        .set_fixes_for_document(fixes, version)
+                        .with_failure_code(ErrorCode::InternalError)?;
                 }
                 Command::OrganizeImports => {
                     let fixes = super::code_action_resolve::organize_imports_edit(

crates/ruff_server/src/server/api/requests/format.rs

@@ -1,9 +1,13 @@
 use crate::edit::{Replacement, ToRangeExt};
+use crate::fix::Fixes;
 use crate::server::api::LSPResult;
 use crate::server::{client::Notifier, Result};
 use crate::session::DocumentSnapshot;
+use crate::{PositionEncoding, TextDocument};
 use lsp_types::{self as types, request as req};
+use ruff_python_ast::PySourceType;
 use ruff_source_file::LineIndex;
+use ruff_workspace::FormatterSettings;
 use types::TextEdit;
 
 pub(crate) struct Format;
@@ -23,25 +27,73 @@ impl super::BackgroundDocumentRequestHandler for Format {
     }
 }
 
+/// Formats either a full text document or each individual cell in a single notebook document.
+pub(super) fn format_full_document(snapshot: &DocumentSnapshot) -> Result<Fixes> {
+    let mut fixes = Fixes::default();
+
+    if let Some(notebook) = snapshot.query().as_notebook() {
+        for (url, text_document) in notebook
+            .urls()
+            .map(|url| (url.clone(), notebook.cell_document_by_uri(url).unwrap()))
+        {
+            if let Some(changes) = format_text_document(
+                text_document,
+                snapshot.query().source_type(),
+                snapshot.query().settings().formatter(),
+                snapshot.encoding(),
+                true,
+            )? {
+                fixes.insert(url, changes);
+            }
+        }
+    } else {
+        if let Some(changes) = format_text_document(
+            snapshot.query().as_single_document().unwrap(),
+            snapshot.query().source_type(),
+            snapshot.query().settings().formatter(),
+            snapshot.encoding(),
+            false,
+        )? {
+            fixes.insert(snapshot.query().make_key().into_url(), changes);
+        }
+    }
+
+    Ok(fixes)
+}
+
+/// Formats either a full text document or an specific notebook cell. If the query within the snapshot is a notebook document
+/// with no selected cell, this will throw an error.
 pub(super) fn format_document(snapshot: &DocumentSnapshot) -> Result<super::FormatResponse> {
-    let doc = snapshot
+    let text_document = snapshot
         .query()
         .as_single_document()
         .expect("format should only be called on text documents or notebook cells");
-    let source = doc.contents();
-    let mut formatted = crate::format::format(
-        doc,
+    format_text_document(
+        text_document,
         snapshot.query().source_type(),
         snapshot.query().settings().formatter(),
+        snapshot.encoding(),
+        snapshot.query().as_notebook().is_some(),
     )
-    .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+}
+
+fn format_text_document(
+    text_document: &TextDocument,
+    source_type: PySourceType,
+    formatter_settings: &FormatterSettings,
+    encoding: PositionEncoding,
+    is_notebook: bool,
+) -> Result<super::FormatResponse> {
+    let source = text_document.contents();
+    let mut formatted = crate::format::format(text_document, source_type, formatter_settings)
+        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
     // fast path - if the code is the same, return early
     if formatted == source {
         return Ok(None);
     }
 
     // special case - avoid adding a newline to a notebook cell if it didn't already exist
-    if snapshot.query().as_notebook().is_some() {
+    if is_notebook {
         let mut trimmed = formatted.as_str();
         if !source.ends_with("\r\n") {
             trimmed = trimmed.trim_end_matches("\r\n");
@@ -57,7 +109,7 @@ pub(super) fn format_document(snapshot: &DocumentSnapshot) -> Result<super::Form
 
     let formatted_index: LineIndex = LineIndex::from_source_text(&formatted);
 
-    let unformatted_index = doc.index();
+    let unformatted_index = text_document.index();
 
     let Replacement {
         source_range,
@@ -70,7 +122,7 @@ pub(super) fn format_document(snapshot: &DocumentSnapshot) -> Result<super::Form
     );
 
     Ok(Some(vec![TextEdit {
-        range: source_range.to_range(source, unformatted_index, snapshot.encoding()),
+        range: source_range.to_range(source, unformatted_index, encoding),
         new_text: formatted[formatted_range].to_owned(),
     }]))
 }