From e244d75df1ea3622dfc9341d0d965fbcae02dfd7 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 28 Mar 2025 14:23:13 -0400 Subject: [PATCH 1/4] [red-knot] Incorporate recent ruff server improvements into red knot's LSP --- crates/red_knot_server/src/edit/notebook.rs | 137 ++++++++++++++++-- .../red_knot_server/src/edit/text_document.rs | 109 +++++++++++++- crates/red_knot_server/src/server/api.rs | 6 +- .../red_knot_server/src/server/connection.rs | 37 ++++- 4 files changed, 262 insertions(+), 27 deletions(-) diff --git a/crates/red_knot_server/src/edit/notebook.rs b/crates/red_knot_server/src/edit/notebook.rs index eb716972a94eb..452263264f8fe 100644 --- a/crates/red_knot_server/src/edit/notebook.rs +++ b/crates/red_knot_server/src/edit/notebook.rs @@ -136,17 +136,15 @@ impl NotebookDocument { // provide the actual contents of the cells, so we'll initialize them with empty // contents. for cell in structure.array.cells.into_iter().flatten().rev() { - if let Some(text_document) = deleted_cells.remove(&cell.document) { - let version = text_document.version(); - self.cells.push(NotebookCell::new( - cell, - text_document.into_contents(), - version, - )); - } else { - self.cells - .insert(start, NotebookCell::new(cell, String::new(), 0)); - } + let (content, version) = + if let Some(text_document) = deleted_cells.remove(&cell.document) { + let version = text_document.version(); + (text_document.into_contents(), version) + } else { + (String::new(), 0) + }; + self.cells + .insert(start, NotebookCell::new(cell, content, version)); } // Third, register the new cells in the index and update existing ones that came @@ -200,6 +198,11 @@ impl NotebookDocument { self.version } + /// Get the URI for a cell by its index within the cell array. + pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> { + self.cells.get(index).map(|cell| &cell.url) + } + /// Get the text document representing the contents of a cell by the cell URI. pub(crate) fn cell_document_by_uri(&self, uri: &lsp_types::Url) -> Option<&TextDocument> { self.cells @@ -238,3 +241,115 @@ impl NotebookCell { } } } + +#[cfg(test)] +mod tests { + use super::NotebookDocument; + + enum TestCellContent { + #[allow(dead_code)] + Markup(String), + Code(String), + } + + fn create_test_url(index: usize) -> lsp_types::Url { + lsp_types::Url::parse(&format!("cell:/test.ipynb#{index}")).unwrap() + } + + fn create_test_notebook(test_cells: Vec) -> NotebookDocument { + let mut cells = Vec::with_capacity(test_cells.len()); + let mut cell_documents = Vec::with_capacity(test_cells.len()); + + for (index, test_cell) in test_cells.into_iter().enumerate() { + let url = create_test_url(index); + match test_cell { + TestCellContent::Markup(content) => { + cells.push(lsp_types::NotebookCell { + kind: lsp_types::NotebookCellKind::Markup, + document: url.clone(), + metadata: None, + execution_summary: None, + }); + cell_documents.push(lsp_types::TextDocumentItem { + uri: url, + language_id: "markdown".to_owned(), + version: 0, + text: content, + }); + } + TestCellContent::Code(content) => { + cells.push(lsp_types::NotebookCell { + kind: lsp_types::NotebookCellKind::Code, + document: url.clone(), + metadata: None, + execution_summary: None, + }); + cell_documents.push(lsp_types::TextDocumentItem { + uri: url, + language_id: "python".to_owned(), + version: 0, + text: content, + }); + } + } + } + + NotebookDocument::new(0, cells, serde_json::Map::default(), cell_documents).unwrap() + } + + /// This test case checks that for a notebook with three code cells, when the client sends a + /// change request to swap the first two cells, the notebook document is updated correctly. + /// + /// The swap operation as a change request is represented as deleting the first two cells and + /// adding them back in the reverse order. + #[test] + fn swap_cells() { + let mut notebook = create_test_notebook(vec![ + TestCellContent::Code("cell = 0".to_owned()), + TestCellContent::Code("cell = 1".to_owned()), + TestCellContent::Code("cell = 2".to_owned()), + ]); + + notebook + .update( + Some(lsp_types::NotebookDocumentCellChange { + structure: Some(lsp_types::NotebookDocumentCellChangeStructure { + array: lsp_types::NotebookCellArrayChange { + start: 0, + delete_count: 2, + cells: Some(vec![ + lsp_types::NotebookCell { + kind: lsp_types::NotebookCellKind::Code, + document: create_test_url(1), + metadata: None, + execution_summary: None, + }, + lsp_types::NotebookCell { + kind: lsp_types::NotebookCellKind::Code, + document: create_test_url(0), + metadata: None, + execution_summary: None, + }, + ]), + }, + did_open: None, + did_close: None, + }), + data: None, + text_content: None, + }), + None, + 1, + crate::PositionEncoding::default(), + ) + .unwrap(); + + assert_eq!( + notebook.make_ruff_notebook().source_code(), + "cell = 1 +cell = 0 +cell = 2 +" + ); + } +} diff --git a/crates/red_knot_server/src/edit/text_document.rs b/crates/red_knot_server/src/edit/text_document.rs index 1d5d496b5bb48..77377ca937984 100644 --- a/crates/red_knot_server/src/edit/text_document.rs +++ b/crates/red_knot_server/src/edit/text_document.rs @@ -20,6 +20,23 @@ pub struct TextDocument { /// The latest version of the document, set by the LSP client. The server will panic in /// debug mode if we attempt to update the document with an 'older' version. version: DocumentVersion, + /// The language ID of the document as provided by the client. + language_id: Option, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum LanguageId { + Python, + Other, +} + +impl From<&str> for LanguageId { + fn from(language_id: &str) -> Self { + match language_id { + "python" => Self::Python, + _ => Self::Other, + } + } } impl TextDocument { @@ -29,9 +46,16 @@ impl TextDocument { contents, index, version, + language_id: None, } } + #[must_use] + pub fn with_language_id(mut self, language_id: &str) -> Self { + self.language_id = Some(LanguageId::from(language_id)); + self + } + pub fn into_contents(self) -> String { self.contents } @@ -48,6 +72,10 @@ impl TextDocument { self.version } + pub fn language_id(&self) -> Option { + self.language_id + } + pub fn apply_changes( &mut self, changes: Vec, @@ -66,7 +94,6 @@ impl TextDocument { return; } - let old_contents = self.contents().to_string(); let mut new_contents = self.contents().to_string(); let mut active_index = self.index().clone(); @@ -87,15 +114,11 @@ impl TextDocument { new_contents = change; } - if new_contents != old_contents { - active_index = LineIndex::from_source_text(&new_contents); - } + active_index = LineIndex::from_source_text(&new_contents); } self.modify_with_manual_index(|contents, version, index| { - if contents != &new_contents { - *index = active_index; - } + *index = active_index; *contents = new_contents; *version = new_version; }); @@ -125,3 +148,75 @@ impl TextDocument { debug_assert!(self.version >= old_version); } } + +#[cfg(test)] +mod tests { + use crate::{PositionEncoding, TextDocument}; + use lsp_types::{Position, TextDocumentContentChangeEvent}; + + #[test] + fn redo_edit() { + let mut document = TextDocument::new( + r#"""" +测试comment +一些测试内容 +""" +import click + + +@click.group() +def interface(): + pas +"# + .to_string(), + 0, + ); + + // Add an `s`, remove it again (back to the original code), and then re-add the `s` + document.apply_changes( + vec![ + TextDocumentContentChangeEvent { + range: Some(lsp_types::Range::new( + Position::new(9, 7), + Position::new(9, 7), + )), + range_length: Some(0), + text: "s".to_string(), + }, + TextDocumentContentChangeEvent { + range: Some(lsp_types::Range::new( + Position::new(9, 7), + Position::new(9, 8), + )), + range_length: Some(1), + text: String::new(), + }, + TextDocumentContentChangeEvent { + range: Some(lsp_types::Range::new( + Position::new(9, 7), + Position::new(9, 7), + )), + range_length: Some(0), + text: "s".to_string(), + }, + ], + 1, + PositionEncoding::UTF16, + ); + + assert_eq!( + &document.contents, + r#"""" +测试comment +一些测试内容 +""" +import click + + +@click.group() +def interface(): + pass +"# + ); + } +} diff --git a/crates/red_knot_server/src/server/api.rs b/crates/red_knot_server/src/server/api.rs index 9638c14e8c181..5270feca4cfac 100644 --- a/crates/red_knot_server/src/server/api.rs +++ b/crates/red_knot_server/src/server/api.rs @@ -69,6 +69,7 @@ fn _local_request_task<'a, R: traits::SyncRequestHandler>( ) -> super::Result> { let (id, params) = cast_request::(req)?; Ok(Task::local(|session, notifier, requester, responder| { + let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered(); let result = R::run(session, notifier, requester, params); respond::(id, result, &responder); })) @@ -98,6 +99,7 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>( }; Box::new(move |notifier, responder| { + let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered(); let result = R::run_with_snapshot(snapshot, db, notifier, params); respond::(id, result, &responder); }) @@ -109,6 +111,7 @@ fn local_notification_task<'a, N: traits::SyncNotificationHandler>( ) -> super::Result> { let (id, params) = cast_notification::(notif)?; Ok(Task::local(move |session, notifier, requester, _| { + let _span = tracing::trace_span!("notification", method = N::METHOD).entered(); if let Err(err) = N::run(session, notifier, requester, params) { tracing::error!("An error occurred while running {id}: {err}"); show_err_msg!("Ruff encountered a problem. Check the logs for more details."); @@ -128,6 +131,7 @@ fn background_notification_thread<'a, N: traits::BackgroundDocumentNotificationH return Box::new(|_, _| {}); }; Box::new(move |notifier, _| { + let _span = tracing::trace_span!("notification", method = N::METHOD).entered(); if let Err(err) = N::run_with_snapshot(snapshot, notifier, params) { tracing::error!("An error occurred while running {id}: {err}"); show_err_msg!("Ruff encountered a problem. Check the logs for more details."); @@ -174,7 +178,7 @@ fn respond( Req: traits::RequestHandler, { if let Err(err) = &result { - tracing::error!("An error occurred with result ID {id}: {err}"); + tracing::error!("An error occurred with request ID {id}: {err}"); show_err_msg!("Ruff encountered a problem. Check the logs for more details."); } if let Err(err) = responder.respond(id, result) { diff --git a/crates/red_knot_server/src/server/connection.rs b/crates/red_knot_server/src/server/connection.rs index c04567c57ae84..5993023be8303 100644 --- a/crates/red_knot_server/src/server/connection.rs +++ b/crates/red_knot_server/src/server/connection.rs @@ -91,19 +91,40 @@ impl Connection { self.sender .send(lsp::Response::new_ok(id.clone(), ()).into())?; tracing::info!("Shutdown request received. Waiting for an exit notification..."); - match self.receiver.recv_timeout(std::time::Duration::from_secs(30))? { - lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => { - tracing::info!("Exit notification received. Server shutting down..."); - Ok(true) - }, - message => anyhow::bail!("Server received unexpected message {message:?} while waiting for exit notification") + + loop { + match &self + .receiver + .recv_timeout(std::time::Duration::from_secs(30))? + { + lsp::Message::Notification(lsp::Notification { method, .. }) + if method == lsp_types::notification::Exit::METHOD => + { + tracing::info!("Exit notification received. Server shutting down..."); + return Ok(true); + } + lsp::Message::Request(lsp::Request { id, method, .. }) => { + tracing::warn!( + "Server received unexpected request {method} ({id}) while waiting for exit notification", + ); + self.sender.send(lsp::Message::Response(lsp::Response::new_err( + id.clone(), + lsp::ErrorCode::InvalidRequest as i32, + "Server received unexpected request while waiting for exit notification".to_string(), + )))?; + } + message => { + tracing::warn!( + "Server received unexpected message while waiting for exit notification: {message:?}" + ); + } + } } } lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => { - tracing::error!("Server received an exit notification before a shutdown request was sent. Exiting..."); - Ok(true) + anyhow::bail!("Server received an exit notification before a shutdown request was sent. Exiting..."); } _ => Ok(false), } From 0dca2e29dd577fe8bd8961386e99b4fd7ebd8725 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 28 Mar 2025 14:28:30 -0400 Subject: [PATCH 2/4] Change logging for not implemented `setTrace` handler to trace --- crates/red_knot_server/src/server/api.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/red_knot_server/src/server/api.rs b/crates/red_knot_server/src/server/api.rs index 5270feca4cfac..6a49f7758abf5 100644 --- a/crates/red_knot_server/src/server/api.rs +++ b/crates/red_knot_server/src/server/api.rs @@ -1,8 +1,8 @@ -use lsp_server as server; - use crate::server::schedule::Task; use crate::session::Session; use crate::system::{url_to_any_system_path, AnySystemPath}; +use lsp_server as server; +use lsp_types::notification::Notification; mod diagnostics; mod notifications; @@ -52,6 +52,11 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> { notification::DidCloseNotebookHandler::METHOD => { local_notification_task::(notif) } + lsp_types::notification::SetTrace::METHOD => { + tracing::trace!("Ignoring `setTrace` notification"); + return Task::nothing(); + }, + method => { tracing::warn!("Received notification {method} which does not have a handler."); return Task::nothing(); From 1a37dea4860507fb9f19e3b078770847c5bf17e0 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 28 Mar 2025 14:32:55 -0400 Subject: [PATCH 3/4] Set language id --- .../src/server/api/notifications/did_open.rs | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/crates/red_knot_server/src/server/api/notifications/did_open.rs b/crates/red_knot_server/src/server/api/notifications/did_open.rs index a42ea1f3ba9d4..2530640bf6231 100644 --- a/crates/red_knot_server/src/server/api/notifications/did_open.rs +++ b/crates/red_knot_server/src/server/api/notifications/did_open.rs @@ -1,5 +1,5 @@ use lsp_types::notification::DidOpenTextDocument; -use lsp_types::DidOpenTextDocumentParams; +use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem}; use red_knot_project::watch::ChangeEvent; use ruff_db::Db; @@ -22,14 +22,22 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler { session: &mut Session, _notifier: Notifier, _requester: &mut Requester, - params: DidOpenTextDocumentParams, + DidOpenTextDocumentParams { + text_document: + TextDocumentItem { + uri, + text, + version, + language_id, + }, + }: DidOpenTextDocumentParams, ) -> Result<()> { - let Ok(path) = url_to_any_system_path(¶ms.text_document.uri) else { + let Ok(path) = url_to_any_system_path(&uri) else { return Ok(()); }; - let document = TextDocument::new(params.text_document.text, params.text_document.version); - session.open_text_document(params.text_document.uri, document); + let document = TextDocument::new(text, version).with_language_id(&language_id); + session.open_text_document(uri, document); match path { AnySystemPath::System(path) => { From fb0c5f40cdf35370bd6aed4c69af46752e945e27 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 28 Mar 2025 14:34:47 -0400 Subject: [PATCH 4/4] Improve Astral logo in dark mode --- playground/shared/src/Header.tsx | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/playground/shared/src/Header.tsx b/playground/shared/src/Header.tsx index 1502db86ad11b..b1dedf064a57a 100644 --- a/playground/shared/src/Header.tsx +++ b/playground/shared/src/Header.tsx @@ -103,32 +103,20 @@ function Logo({ fillRule="evenodd" clipRule="evenodd" d="M431.998 9.98526C431.998 4.47055 436.469 0 441.984 0H522.013C527.528 0 531.998 4.47056 531.998 9.98526V100H485.998V70H477.998V100H431.998V9.98526ZM493.998 46V38H469.998V46H493.998Z" - fill="#30173D" /> - - + + - + ); }