From 6e9ff445fd8559972b423370de20563a9c2db8d4 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Fri, 10 Jan 2025 18:41:56 +0530 Subject: [PATCH] Insert the cells from the `start` position (#15398) ## Summary The cause of this bug is from https://github.com/astral-sh/ruff/pull/12575 which was itself a bug fix but the fix wasn't completely correct. fixes: #14768 fixes: https://github.com/astral-sh/ruff-vscode/issues/644 ## Test Plan Consider the following three cells: 1. ```python class Foo: def __init__(self): self.x = 1 def __str__(self): return f"Foo({self.x})" ``` 2. ```python def hello(): print("hello world") ``` 3. ```python y = 1 ``` The test case is moving cell 2 to the top i.e., cell 2 goes to position 1 and cell 1 goes to position 2. Before this fix, it can be seen that the cells were pushed at the end of the vector: ``` 12.643269917s INFO ruff:main ruff_server::edit::notebook: Before update: [ NotebookCell { document: TextDocument { contents: "class Foo:\n def __init__(self):\n self.x = 1\n\n def __str__(self):\n return f\"Foo({self.x})\"", }, }, NotebookCell { document: TextDocument { contents: "def hello():\n print(\"hello world\")", }, }, NotebookCell { document: TextDocument { contents: "y = 1", }, }, ] 12.643777667s INFO ruff:main ruff_server::edit::notebook: After update: [ NotebookCell { document: TextDocument { contents: "y = 1", }, }, NotebookCell { document: TextDocument { contents: "class Foo:\n def __init__(self):\n self.x = 1\n\n def __str__(self):\n return f\"Foo({self.x})\"", }, }, NotebookCell { document: TextDocument { contents: "def hello():\n print(\"hello world\")", }, }, ] ``` After the fix in this PR, it can be seen that the cells are being pushed at the correct `start` index: ``` 6.520570917s INFO ruff:main ruff_server::edit::notebook: Before update: [ NotebookCell { document: TextDocument { contents: "class Foo:\n def __init__(self):\n self.x = 1\n\n def __str__(self):\n return f\"Foo({self.x})\"", }, }, NotebookCell { document: TextDocument { contents: "def hello():\n print(\"hello world\")", }, }, NotebookCell { document: TextDocument { contents: "y = 1", }, }, ] 6.521084792s INFO ruff:main ruff_server::edit::notebook: After update: [ NotebookCell { document: TextDocument { contents: "def hello():\n print(\"hello world\")", }, }, NotebookCell { document: TextDocument { contents: "class Foo:\n def __init__(self):\n self.x = 1\n\n def __str__(self):\n return f\"Foo({self.x})\"", }, }, NotebookCell { document: TextDocument { contents: "y = 1", }, }, ] ``` --- crates/ruff_server/src/edit/notebook.rs | 132 ++++++++++++++++++++++-- 1 file changed, 121 insertions(+), 11 deletions(-) diff --git a/crates/ruff_server/src/edit/notebook.rs b/crates/ruff_server/src/edit/notebook.rs index 7e0c65e9174f3..452263264f8fe 100644 --- a/crates/ruff_server/src/edit/notebook.rs +++ b/crates/ruff_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 @@ -243,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 +" + ); + } +}