From 8cb7d556e08f232f01460604a58a46b0bdcd9dab Mon Sep 17 00:00:00 2001 From: Jane Lewis Date: Wed, 15 May 2024 21:38:37 -0700 Subject: [PATCH] Address additional suggestions --- crates/ruff_server/src/fix.rs | 22 +++---- crates/ruff_server/src/lint.rs | 64 ++++++++---------- .../server/api/notifications/did_change.rs | 7 +- .../api/notifications/did_change_notebook.rs | 4 +- .../api/notifications/did_close_notebook.rs | 4 +- crates/ruff_server/src/session.rs | 9 ++- crates/ruff_server/src/session/index.rs | 65 ++++++++++--------- 7 files changed, 87 insertions(+), 88 deletions(-) diff --git a/crates/ruff_server/src/fix.rs b/crates/ruff_server/src/fix.rs index 664a6ba0f0f33..e1fa2e281db53 100644 --- a/crates/ruff_server/src/fix.rs +++ b/crates/ruff_server/src/fix.rs @@ -70,6 +70,13 @@ 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"); }; @@ -77,19 +84,8 @@ pub(crate) fn fix_all( 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); diff --git a/crates/ruff_server/src/lint.rs b/crates/ruff_server/src/lint.rs index b990120cdf39a..91aa2424e1d2d 100644 --- a/crates/ruff_server/src/lint.rs +++ b/crates/ruff_server/src/lint.rs @@ -15,7 +15,7 @@ use ruff_python_codegen::Stylist; use ruff_python_index::Indexer; 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}; @@ -127,20 +127,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() @@ -208,42 +211,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 { @@ -293,6 +266,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 diff --git a/crates/ruff_server/src/server/api/notifications/did_change.rs b/crates/ruff_server/src/server/api/notifications/did_change.rs index 5e3dc392b16c2..2d4653363a9da 100644 --- a/crates/ruff_server/src/server/api/notifications/did_change.rs +++ b/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 { diff --git a/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs b/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs index 48d688a54826f..b0ddb4747c35e 100644 --- a/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs +++ b/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)?; diff --git a/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs b/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs index 8b01651e12270..a66e3126e5cb9 100644 --- a/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs +++ b/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) diff --git a/crates/ruff_server/src/session.rs b/crates/ruff_server/src/session.rs index 57585445a2261..3c271bb6fda9b 100644 --- a/crates/ruff_server/src/session.rs +++ b/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 { + 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 { - 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), diff --git a/crates/ruff_server/src/session/index.rs b/crates/ruff_server/src/session/index.rs index 167762986e9b9..9304cba1fccdc 100644 --- a/crates/ruff_server/src/session/index.rs +++ b/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, }, 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, file_path: PathBuf, notebook: Arc, @@ -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 { 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,10 +136,6 @@ 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, @@ -153,6 +143,10 @@ impl Index { .. }) = 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),