ruff server: Support Jupyter Notebook (*.ipynb) files (#11206)

## Summary

Closes https://github.com/astral-sh/ruff/issues/10858.

`ruff server` now supports `*.ipynb` (aka Jupyter Notebook) files.
Extensive internal changes have been made to facilitate this, which I've
done some work to contextualize with documentation and an pre-review
that highlights notable sections of the code.

`*.ipynb` cells should behave similarly to `*.py` documents, with one
major exception. The format command `ruff.applyFormat` will only apply
to the currently selected notebook cell - if you want to format an
entire notebook document, use `Format Notebook` from the VS Code context
menu.

## Test Plan

The VS Code extension does not yet have Jupyter Notebook support
enabled, so you'll first need to enable it manually. To do this,
checkout the `pre-release` branch and modify `src/common/server.ts` as
follows:

Before:
![Screenshot 2024-05-13 at 10 59
06 PM](c6a3c604-c405-4968-b8a2-5d670de89172)

After:
![Screenshot 2024-05-13 at 10 58
24 PM](94ab2e3d-0609-448d-9c8c-cd07c69a513b)

I recommend testing this PR with large, complicated notebook files. I
used notebook files from [this popular
repository](https://github.com/jakevdp/PythonDataScienceHandbook/tree/master/notebooks)
in my preliminary testing.

The main thing to test is ensuring that notebook cells behave the same
as Python documents, besides the aforementioned issue with
`ruff.applyFormat`. You should also test adding and deleting cells (in
particular, deleting all the code cells and ensure that doesn't break
anything), changing the kind of a cell (i.e. from markup -> code or vice
versa), and creating a new notebook file from scratch. Finally, you
should also test that source actions work as expected (and across the
entire notebook).

Note: `ruff.applyAutofix` and `ruff.applyOrganizeImports` are currently
broken for notebook files, and I suspect it has something to do with
https://github.com/astral-sh/ruff/issues/11248. Once this is fixed, I
will update the test plan accordingly.

---------

Co-authored-by: nolan <nolan.king90@gmail.com>

crates/ruff_server/src/session.rs

@@ -1,24 +1,26 @@
 //! Data model, state management, and configuration resolution.
 
 mod capabilities;
+mod index;
 mod settings;
-mod workspace;
 
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use anyhow::anyhow;
-use lsp_types::{ClientCapabilities, Url};
+use lsp_types::{ClientCapabilities, NotebookDocumentCellChange, Url};
 
-use crate::edit::DocumentVersion;
-use crate::PositionEncoding;
+use crate::edit::{DocumentKey, DocumentVersion, NotebookDocument};
+use crate::{PositionEncoding, TextDocument};
 
 pub(crate) use self::capabilities::ResolvedClientCapabilities;
+pub(crate) use self::index::DocumentQuery;
 pub(crate) use self::settings::{AllSettings, ClientSettings};
 
 /// The global state for the LSP
 pub(crate) struct Session {
-    /// Workspace folders in the current session, which contain the state of all open files.
-    workspaces: workspace::Workspaces,
+    /// Used to retrieve information about open documents and settings.
+    index: index::Index,
     /// The global position encoding, negotiated during LSP initialization.
     position_encoding: PositionEncoding,
     /// Global settings provided by the client.
@@ -32,9 +34,8 @@ pub(crate) struct Session {
 pub(crate) struct DocumentSnapshot {
     resolved_client_capabilities: Arc<ResolvedClientCapabilities>,
     client_settings: settings::ResolvedClientSettings,
-    document_ref: workspace::DocumentRef,
+    document_ref: index::DocumentQuery,
     position_encoding: PositionEncoding,
-    url: Url,
 }
 
 impl Session {
@@ -42,58 +43,100 @@ impl Session {
         client_capabilities: &ClientCapabilities,
         position_encoding: PositionEncoding,
         global_settings: ClientSettings,
-        workspaces: Vec<(Url, ClientSettings)>,
-    ) -> crate::Result<Self> {
-        Ok(Self {
+        workspace_folders: Vec<(PathBuf, ClientSettings)>,
+    ) -> Self {
+        Self {
             position_encoding,
-            workspaces: workspace::Workspaces::new(workspaces, &global_settings)?,
+            index: index::Index::new(workspace_folders, &global_settings),
             global_settings,
             resolved_client_capabilities: Arc::new(ResolvedClientCapabilities::new(
                 client_capabilities,
             )),
-        })
+        }
     }
 
+    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).ok()?;
         Some(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
-            client_settings: self.workspaces.client_settings(url, &self.global_settings),
-            document_ref: self.workspaces.snapshot(url)?,
+            client_settings: self.index.client_settings(&key, &self.global_settings),
+            document_ref: self.index.make_document_ref(key)?,
             position_encoding: self.position_encoding,
-            url: url.clone(),
         })
     }
 
-    pub(crate) fn open_document(&mut self, url: &Url, contents: String, version: DocumentVersion) {
-        self.workspaces.open(url, contents, version);
-    }
-
-    pub(crate) fn close_document(&mut self, url: &Url) -> crate::Result<()> {
-        self.workspaces.close(url)?;
-        Ok(())
-    }
-
-    pub(crate) fn document_controller(
+    /// Updates a text document at the associated `key`.
+    ///
+    /// The document key must point to a text document, or this will throw an error.
+    pub(crate) fn update_text_document(
         &mut self,
-        url: &Url,
-    ) -> crate::Result<&mut workspace::DocumentController> {
-        self.workspaces
-            .controller(url)
-            .ok_or_else(|| anyhow!("Tried to open unavailable document `{url}`"))
+        key: &DocumentKey,
+        content_changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
+        new_version: DocumentVersion,
+    ) -> crate::Result<()> {
+        let encoding = self.encoding();
+
+        self.index
+            .update_text_document(key, content_changes, new_version, encoding)
     }
 
-    pub(crate) fn reload_settings(&mut self, url: &Url) -> crate::Result<()> {
-        self.workspaces.reload_settings(url)
+    /// 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, or this will
+    /// throw an error.
+    pub(crate) fn update_notebook_document(
+        &mut self,
+        key: &DocumentKey,
+        cells: Option<NotebookDocumentCellChange>,
+        metadata: Option<serde_json::Map<String, serde_json::Value>>,
+        version: DocumentVersion,
+    ) -> crate::Result<()> {
+        let encoding = self.encoding();
+        self.index
+            .update_notebook_document(key, cells, metadata, version, encoding)
     }
 
-    pub(crate) fn open_workspace_folder(&mut self, url: &Url) -> crate::Result<()> {
-        self.workspaces
-            .open_workspace_folder(url, &self.global_settings)?;
+    /// Registers a notebook document at the provided `path`.
+    /// If a document is already open here, it will be overwritten.
+    pub(crate) fn open_notebook_document(&mut self, path: PathBuf, document: NotebookDocument) {
+        self.index.open_notebook_document(path, document);
+    }
+
+    /// Registers a text document at the provided `path`.
+    /// If a document is already open here, it will be overwritten.
+    pub(crate) fn open_text_document(&mut self, path: PathBuf, document: TextDocument) {
+        self.index.open_text_document(path, document);
+    }
+
+    /// De-registers a document, specified by its key.
+    /// Calling this multiple times for the same document is a logic error.
+    pub(crate) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
+        self.index.close_document(key)?;
         Ok(())
     }
 
-    pub(crate) fn close_workspace_folder(&mut self, url: &Url) -> crate::Result<()> {
-        self.workspaces.close_workspace_folder(url)?;
+    /// Reloads the settings index
+    pub(crate) fn reload_settings(&mut self, changed_path: &PathBuf) {
+        self.index.reload_settings(changed_path);
+    }
+
+    /// Open a workspace folder at the given `path`.
+    pub(crate) fn open_workspace_folder(&mut self, path: PathBuf) {
+        self.index
+            .open_workspace_folder(path, &self.global_settings);
+    }
+
+    /// Close a workspace folder at the given `path`.
+    pub(crate) fn close_workspace_folder(&mut self, path: &PathBuf) -> crate::Result<()> {
+        self.index.close_workspace_folder(path)?;
         Ok(())
     }
 
@@ -107,10 +150,6 @@ impl Session {
 }
 
 impl DocumentSnapshot {
-    pub(crate) fn settings(&self) -> &workspace::RuffSettings {
-        self.document().settings()
-    }
-
     pub(crate) fn resolved_client_capabilities(&self) -> &ResolvedClientCapabilities {
         &self.resolved_client_capabilities
     }
@@ -119,15 +158,11 @@ impl DocumentSnapshot {
         &self.client_settings
     }
 
-    pub(crate) fn document(&self) -> &workspace::DocumentRef {
+    pub(crate) fn query(&self) -> &index::DocumentQuery {
         &self.document_ref
     }
 
     pub(crate) fn encoding(&self) -> PositionEncoding {
         self.position_encoding
     }
-
-    pub(crate) fn url(&self) -> &Url {
-        &self.url
-    }
 }