Implement support for Jupyter Notebooks in ruff server

crates/ruff_notebook/src/cell.rs

@@ -23,7 +23,7 @@ impl fmt::Display for SourceValue {
 
 impl Cell {
     /// Return the [`SourceValue`] of the cell.
-    pub(crate) fn source(&self) -> &SourceValue {
+    pub fn source(&self) -> &SourceValue {
         match self {
             Cell::Code(cell) => &cell.source,
             Cell::Markdown(cell) => &cell.source,

crates/ruff_notebook/src/notebook.rs

@@ -1,5 +1,5 @@
 use std::cmp::Ordering;
-use std::collections::HashSet;
+use std::collections::{BTreeMap, HashSet};
 use std::fs::File;
 use std::io::{BufReader, Cursor, Read, Seek, SeekFrom, Write};
 use std::path::Path;
@@ -51,6 +51,15 @@ pub enum NotebookError {
     InvalidFormat(i64),
 }
 
+pub enum NotebookSource {
+    File(Notebook),
+    Server {
+        source_code: String,
+        index: OnceCell<NotebookIndex>,
+        cell_offsets: CellOffsets,
+    },
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Notebook {
     /// Python source code of the notebook.
@@ -85,6 +94,24 @@ impl Notebook {
         Self::from_reader(Cursor::new(source_code))
     }
 
+    pub fn from_cells(cells: Vec<Cell>) -> Result<Self, NotebookError> {
+        let raw_notebook = RawNotebook {
+            cells,
+            metadata: crate::RawNotebookMetadata {
+                authors: None,
+                kernelspec: None,
+                language_info: None,
+                orig_nbformat: None,
+                title: None,
+                extra: BTreeMap::default(),
+            },
+            nbformat: 4,
+            nbformat_minor: 5,
+        };
+
+        Self::from_raw(raw_notebook, false)
+    }
+
     /// Read a Jupyter Notebook from a [`Read`] implementer.
     ///
     /// See also the black implementation
@@ -113,7 +140,13 @@ impl Notebook {
                 });
             }
         };
+        Self::from_raw(raw_notebook, trailing_newline)
+    }
 
+    fn from_raw(
+        mut raw_notebook: RawNotebook,
+        trailing_newline: bool,
+    ) -> Result<Self, NotebookError> {
         // v4 is what everybody uses
         if raw_notebook.nbformat != 4 {
             // bail because we should have already failed at the json schema stage

crates/ruff_server/Cargo.toml

@@ -22,6 +22,7 @@ ruff_python_formatter = { path = "../ruff_python_formatter" }
 ruff_python_index = { path = "../ruff_python_index" }
 ruff_python_parser = { path = "../ruff_python_parser" }
 ruff_source_file = { path = "../ruff_source_file" }
+ruff_notebook = { path = "../ruff_notebook" }
 ruff_text_size = { path = "../ruff_text_size" }
 ruff_workspace = { path = "../ruff_workspace" }
 

crates/ruff_server/src/edit.rs

@@ -1,18 +1,29 @@
 //! Types and utilities for working with text, modifying source files, and `Ruff <-> LSP` type conversion.
 
 mod document;
+mod notebook;
 mod range;
 mod replacement;
 
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    path::{Display, PathBuf},
+};
 
-pub use document::Document;
+use anyhow::anyhow;
 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 replacement::Replacement;
+use ruff_linter::source_kind::SourceKind;
+use ruff_source_file::LineIndex;
 
-use crate::session::ResolvedClientCapabilities;
+use crate::{
+    fix::Fixes,
+    session::{DocumentQuery, ResolvedClientCapabilities},
+};
 
 /// A convenient enumeration for supported text encodings. Can be converted to [`lsp_types::PositionEncodingKind`].
 // Please maintain the order from least to greatest priority for the derived `Ord` impl.
@@ -29,6 +40,108 @@ pub enum PositionEncoding {
     UTF8,
 }
 
+/// A wrapper for a document. Can be either a text document (`.py`)
+/// or a notebook document (`.ipynb`).
+#[derive(Clone)]
+pub(crate) enum Document {
+    Text(TextDocument),
+    Notebook(NotebookDocument),
+}
+
+#[derive(Clone, Debug)]
+pub(crate) enum DocumentKey {
+    File(PathBuf),
+    Cell(lsp_types::Url),
+}
+
+impl Document {
+    pub(crate) fn version(&self) -> DocumentVersion {
+        match self {
+            Self::Notebook(notebook) => notebook.version(),
+            Self::Text(py) => py.version(),
+        }
+    }
+
+    pub(crate) fn make_source_kind(&self) -> SourceKind {
+        match self {
+            Self::Notebook(notebook) => {
+                let notebook = SourceKind::IpyNotebook(notebook.make_ruff_notebook());
+                tracing::info!("{notebook:?}");
+                notebook
+            }
+            Self::Text(text) => SourceKind::Python(text.contents().to_string()),
+        }
+    }
+
+    pub(crate) fn as_notebook(&self) -> Option<&NotebookDocument> {
+        match self {
+            Self::Notebook(notebook) => Some(&notebook),
+            Self::Text(_) => None,
+        }
+    }
+
+    pub(crate) fn as_text(&self) -> Option<&TextDocument> {
+        match self {
+            Self::Text(py) => Some(&py),
+            Self::Notebook(_) => None,
+        }
+    }
+
+    pub(crate) fn as_notebook_mut(&mut self) -> Option<&mut NotebookDocument> {
+        match self {
+            Self::Notebook(ref mut notebook) => Some(notebook),
+            Self::Text(_) => None,
+        }
+    }
+
+    pub(crate) fn as_text_mut(&mut self) -> Option<&mut TextDocument> {
+        match self {
+            Self::Notebook(_) => None,
+            Self::Text(ref mut py) => Some(py),
+        }
+    }
+
+    pub(crate) fn as_notebook_cell(&self, uri: &lsp_types::Url) -> Option<&TextDocument> {
+        self.as_notebook()?.cell_document_by_uri(uri)
+    }
+
+    pub(crate) fn as_notebook_cell_mut(
+        &mut self,
+        uri: &lsp_types::Url,
+    ) -> Option<&mut TextDocument> {
+        self.as_notebook_mut()?.cell_document_by_uri_mut(uri)
+    }
+}
+
+impl DocumentKey {
+    pub(crate) fn from_url(url: &lsp_types::Url) -> Self {
+        if url.scheme() != "file" {
+            return Self::Cell(url.clone());
+        }
+        url.to_file_path()
+            .map(Self::File)
+            .unwrap_or_else(|_| Self::Cell(url.clone()))
+    }
+
+    pub(crate) fn to_url(self) -> lsp_types::Url {
+        match self {
+            Self::Cell(url) => url,
+            Self::File(path) => {
+                lsp_types::Url::from_file_path(path).expect("path should convert to file URL")
+            }
+        }
+    }
+}
+
+impl std::fmt::Display for DocumentKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Cell(url) => url.fmt(f),
+            Self::File(path) => path.display().fmt(f),
+        }
+    }
+}
+
 /// Tracks multi-document edits to eventually merge into a `WorkspaceEdit`.
 /// Compatible with clients that don't support `workspace.workspaceEdit.documentChanges`.
 #[derive(Debug)]
@@ -72,6 +185,26 @@ impl WorkspaceEditTracker {
         }
     }
 
+    pub(crate) fn set_fixes_for_document(
+        &mut self,
+        document: &DocumentQuery,
+        fixes: Fixes,
+    ) -> crate::Result<()> {
+        for (cell, edits) in fixes.into_iter() {
+            let uri = if let Some(notebook) = document.as_notebook() {
+                notebook
+                    .cell_uri_by_index(cell)
+                    .cloned()
+                    .ok_or_else(|| anyhow!("cell index {cell} does not exist"))?
+            } else {
+                document.key().clone().to_url()
+            };
+
+            self.set_edits_for_document(uri, document.version(), edits)?;
+        }
+        Ok(())
+    }
+
     /// Sets the edits made to a specific document. This should only be called
     /// once for each document `uri`, and will fail if this is called for the same `uri`
     /// multiple times.

crates/ruff_server/src/edit/document.rs

@@ -10,7 +10,7 @@ pub(crate) type DocumentVersion = i32;
 /// The state for an individual document in the server. Stays up-to-date
 /// with changes made by the user, including unsaved changes.
 #[derive(Debug, Clone)]
-pub struct Document {
+pub struct TextDocument {
     /// The string contents of the document.
     contents: String,
     /// A computed line index for the document. This should always reflect
@@ -22,7 +22,7 @@ pub struct Document {
     version: DocumentVersion,
 }
 
-impl Document {
+impl TextDocument {
     pub fn new(contents: String, version: DocumentVersion) -> Self {
         let index = LineIndex::from_source_text(&contents);
         Self {