From d0bcf56bd94a9525a818fd593606b4f5a62add7b Mon Sep 17 00:00:00 2001 From: Brent Westbrook <36778786+ntBre@users.noreply.github.com> Date: Mon, 25 Aug 2025 09:20:42 -0400 Subject: [PATCH] Improve diff rendering for notebooks (#20036) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary As noted in a code TODO, our `Diff` rendering code previously didn't have any special handling for notebooks. This was particularly obvious when the diffs were rendered right next to the corresponding diagnostic because the diagnostic used cell-based line numbers, while the diff was still using line numbers from the concatenated source. This PR updates the diff rendering to handle notebooks too. The main improvements shown in the example below are: - Line numbers are now remapped to be relative to their cell - Context lines from other cells are suppressed ``` error[unused-import][*]: `math` imported but unused --> notebook.ipynb:cell 2:2:8 | 1 | # cell 2 2 | import math | ^^^^ 3 | 4 | print('hello world') | help: Remove unused import: `math` ℹ Safe fix 1 1 | # cell 2 2 |-import math 3 2 | 4 3 | print('hello world') ``` I tried a few different approaches here before finally just splitting the notebook into separate text ranges by cell and diffing each one separately. It seems to work and passes all of our tests, but I don't know if it's actually enforced anywhere that a single edit doesn't span cells. Such an edit would silently be dropped right now since it would fail the `contains_range` check. I also feel like I may have overlooked an existing way to partition a file into cells like this. ## Test Plan Existing notebook tests, plus a new one in `ruff_db` --- crates/ruff_db/src/diagnostic/mod.rs | 5 + crates/ruff_db/src/diagnostic/render.rs | 7 + crates/ruff_db/src/diagnostic/render/full.rs | 273 ++++++++++++++---- ...linter__linter__tests__import_sorting.snap | 53 ++-- ...er__linter__tests__ipy_escape_command.snap | 9 +- ...inter__linter__tests__unused_variable.snap | 28 +- ...er__linter__tests__vscode_language_id.snap | 1 + crates/ruff_notebook/src/index.rs | 8 + 8 files changed, 270 insertions(+), 114 deletions(-) diff --git a/crates/ruff_db/src/diagnostic/mod.rs b/crates/ruff_db/src/diagnostic/mod.rs index db6befc9aa..a7e85edaf3 100644 --- a/crates/ruff_db/src/diagnostic/mod.rs +++ b/crates/ruff_db/src/diagnostic/mod.rs @@ -325,6 +325,11 @@ impl Diagnostic { self.inner.fix.as_ref() } + #[cfg(test)] + pub(crate) fn fix_mut(&mut self) -> Option<&mut Fix> { + Arc::make_mut(&mut self.inner).fix.as_mut() + } + /// Set the fix for this diagnostic. pub fn set_fix(&mut self, fix: Fix) { debug_assert!( diff --git a/crates/ruff_db/src/diagnostic/render.rs b/crates/ruff_db/src/diagnostic/render.rs index 2b5230117e..35cf2ff051 100644 --- a/crates/ruff_db/src/diagnostic/render.rs +++ b/crates/ruff_db/src/diagnostic/render.rs @@ -2622,6 +2622,13 @@ watermelon self.config = config; } + /// Show a diff for the fix when rendering. + pub(super) fn show_fix_diff(&mut self, yes: bool) { + let mut config = std::mem::take(&mut self.config); + config = config.show_fix_diff(yes); + self.config = config; + } + /// The lowest fix applicability to show when rendering. pub(super) fn fix_applicability(&mut self, applicability: Applicability) { let mut config = std::mem::take(&mut self.config); diff --git a/crates/ruff_db/src/diagnostic/render/full.rs b/crates/ruff_db/src/diagnostic/render/full.rs index 951b812562..847386a0f5 100644 --- a/crates/ruff_db/src/diagnostic/render/full.rs +++ b/crates/ruff_db/src/diagnostic/render/full.rs @@ -2,12 +2,13 @@ use std::borrow::Cow; use std::num::NonZeroUsize; use anstyle::Style; +use ruff_notebook::NotebookIndex; use similar::{ChangeTag, TextDiff}; use ruff_annotate_snippets::Renderer as AnnotateRenderer; use ruff_diagnostics::{Applicability, Fix}; use ruff_source_file::OneIndexed; -use ruff_text_size::{Ranged, TextRange, TextSize}; +use ruff_text_size::{Ranged, TextLen, TextRange, TextSize}; use crate::diagnostic::render::{FileResolver, Resolved}; use crate::diagnostic::stylesheet::{DiagnosticStylesheet, fmt_styled}; @@ -81,6 +82,7 @@ impl<'a> FullRenderer<'a> { struct Diff<'a> { fix: &'a Fix, diagnostic_source: DiagnosticSource, + notebook_index: Option, stylesheet: &'a DiagnosticStylesheet, } @@ -90,12 +92,11 @@ impl<'a> Diff<'a> { stylesheet: &'a DiagnosticStylesheet, resolver: &'a dyn FileResolver, ) -> Option> { + let file = &diagnostic.primary_span_ref()?.file; Some(Diff { fix: diagnostic.fix()?, - diagnostic_source: diagnostic - .primary_span_ref()? - .file - .diagnostic_source(resolver), + diagnostic_source: file.diagnostic_source(resolver), + notebook_index: resolver.notebook_index(file), stylesheet, }) } @@ -106,19 +107,24 @@ impl std::fmt::Display for Diff<'_> { let source_code = self.diagnostic_source.as_source_code(); let source_text = source_code.text(); - // TODO(dhruvmanila): Add support for Notebook cells once it's user-facing - let mut output = String::with_capacity(source_text.len()); - let mut last_end = TextSize::default(); - - for edit in self.fix.edits() { - output.push_str(source_code.slice(TextRange::new(last_end, edit.start()))); - output.push_str(edit.content().unwrap_or_default()); - last_end = edit.end(); - } - - output.push_str(&source_text[usize::from(last_end)..]); - - let diff = TextDiff::from_lines(source_text, &output); + // Partition the source code into end offsets for each cell. If `self.notebook_index` is + // `None`, indicating a regular script file, all the lines will be in one "cell" under the + // `None` key. + let cells = if let Some(notebook_index) = &self.notebook_index { + let mut last_cell = OneIndexed::MIN; + let mut cells: Vec<(Option, TextSize)> = Vec::new(); + for (row, cell) in notebook_index.iter() { + if cell != last_cell { + let offset = source_code.line_start(row); + cells.push((Some(last_cell), offset)); + last_cell = cell; + } + } + cells.push((Some(last_cell), source_text.text_len())); + cells + } else { + vec![(None, source_text.text_len())] + }; let message = match self.fix.applicability() { // TODO(zanieb): Adjust this messaging once it's user-facing @@ -133,59 +139,97 @@ impl std::fmt::Display for Diff<'_> { // tests, which is the only place these are currently used. writeln!(f, "ℹ {}", fmt_styled(message, self.stylesheet.separator))?; - let (largest_old, largest_new) = diff - .ops() - .last() - .map(|op| (op.old_range().start, op.new_range().start)) - .unwrap_or_default(); + let mut last_end = TextSize::ZERO; + for (cell, offset) in cells { + let range = TextRange::new(last_end, offset); + last_end = offset; + let input = source_code.slice(range); - let digit_with = OneIndexed::from_zero_indexed(largest_new.max(largest_old)).digits(); + let mut output = String::with_capacity(input.len()); + let mut last_end = range.start(); - for (idx, group) in diff.grouped_ops(3).iter().enumerate() { - if idx > 0 { - writeln!(f, "{:-^1$}", "-", 80)?; + let mut applied = 0; + for edit in self.fix.edits() { + if range.contains_range(edit.range()) { + output.push_str(source_code.slice(TextRange::new(last_end, edit.start()))); + output.push_str(edit.content().unwrap_or_default()); + last_end = edit.end(); + applied += 1; + } } - for op in group { - for change in diff.iter_inline_changes(op) { - let sign = match change.tag() { - ChangeTag::Delete => "-", - ChangeTag::Insert => "+", - ChangeTag::Equal => " ", - }; - let line_style = LineStyle::from(change.tag(), self.stylesheet); + // No edits were applied, so there's no need to diff. + if applied == 0 { + continue; + } - let old_index = change.old_index().map(OneIndexed::from_zero_indexed); - let new_index = change.new_index().map(OneIndexed::from_zero_indexed); + output.push_str(&source_text[usize::from(last_end)..usize::from(range.end())]); - write!( - f, - "{} {} |{}", - Line { - index: old_index, - width: digit_with - }, - Line { - index: new_index, - width: digit_with - }, - fmt_styled(line_style.apply_to(sign), self.stylesheet.emphasis), - )?; + let diff = TextDiff::from_lines(input, &output); - for (emphasized, value) in change.iter_strings_lossy() { - let value = show_nonprinting(&value); - if emphasized { - write!( - f, - "{}", - fmt_styled(line_style.apply_to(&value), self.stylesheet.underline) - )?; - } else { - write!(f, "{}", line_style.apply_to(&value))?; + let (largest_old, largest_new) = diff + .ops() + .last() + .map(|op| (op.old_range().start, op.new_range().start)) + .unwrap_or_default(); + + let digit_with = OneIndexed::from_zero_indexed(largest_new.max(largest_old)).digits(); + + if let Some(cell) = cell { + // Room for 2 digits, 2 x 1 space before each digit, 1 space, and 1 `|`. This + // centers the three colons on the pipe. + writeln!(f, "{:>1$} cell {cell}", ":::", 2 * digit_with.get() + 4)?; + } + + for (idx, group) in diff.grouped_ops(3).iter().enumerate() { + if idx > 0 { + writeln!(f, "{:-^1$}", "-", 80)?; + } + for op in group { + for change in diff.iter_inline_changes(op) { + let sign = match change.tag() { + ChangeTag::Delete => "-", + ChangeTag::Insert => "+", + ChangeTag::Equal => " ", + }; + + let line_style = LineStyle::from(change.tag(), self.stylesheet); + + let old_index = change.old_index().map(OneIndexed::from_zero_indexed); + let new_index = change.new_index().map(OneIndexed::from_zero_indexed); + + write!( + f, + "{} {} |{}", + Line { + index: old_index, + width: digit_with, + }, + Line { + index: new_index, + width: digit_with, + }, + fmt_styled(line_style.apply_to(sign), self.stylesheet.emphasis), + )?; + + for (emphasized, value) in change.iter_strings_lossy() { + let value = show_nonprinting(&value); + if emphasized { + write!( + f, + "{}", + fmt_styled( + line_style.apply_to(&value), + self.stylesheet.underline + ) + )?; + } else { + write!(f, "{}", line_style.apply_to(&value))?; + } + } + if change.missing_newline() { + writeln!(f)?; } - } - if change.missing_newline() { - writeln!(f)?; } } } @@ -253,7 +297,7 @@ fn show_nonprinting(s: &str) -> Cow<'_, str> { #[cfg(test)] mod tests { - use ruff_diagnostics::Applicability; + use ruff_diagnostics::{Applicability, Fix}; use ruff_text_size::{TextLen, TextRange, TextSize}; use crate::diagnostic::{ @@ -654,6 +698,107 @@ print() "); } + /// Test that we remap notebook cell line numbers in the diff as well as the main diagnostic. + #[test] + fn notebook_output_with_diff() { + let (mut env, diagnostics) = create_notebook_diagnostics(DiagnosticFormat::Full); + env.show_fix_diff(true); + insta::assert_snapshot!(env.render_diagnostics(&diagnostics), @r" + error[unused-import][*]: `os` imported but unused + --> notebook.ipynb:cell 1:2:8 + | + 1 | # cell 1 + 2 | import os + | ^^ + | + help: Remove unused import: `os` + + ℹ Safe fix + ::: cell 1 + 1 1 | # cell 1 + 2 |-import os + + error[unused-import][*]: `math` imported but unused + --> notebook.ipynb:cell 2:2:8 + | + 1 | # cell 2 + 2 | import math + | ^^^^ + 3 | + 4 | print('hello world') + | + help: Remove unused import: `math` + + ℹ Safe fix + ::: cell 2 + 1 1 | # cell 2 + 2 |-import math + 3 2 | + 4 3 | print('hello world') + + error[unused-variable]: Local variable `x` is assigned to but never used + --> notebook.ipynb:cell 3:4:5 + | + 2 | def foo(): + 3 | print() + 4 | x = 1 + | ^ + | + help: Remove assignment to unused variable `x` + + ℹ Unsafe fix + ::: cell 3 + 1 1 | # cell 3 + 2 2 | def foo(): + 3 3 | print() + 4 |- x = 1 + 5 4 | + "); + } + + #[test] + fn notebook_output_with_diff_spanning_cells() { + let (mut env, mut diagnostics) = create_notebook_diagnostics(DiagnosticFormat::Full); + env.show_fix_diff(true); + + // Move all of the edits from the later diagnostics to the first diagnostic to simulate a + // single diagnostic with edits in different cells. + let mut diagnostic = diagnostics.swap_remove(0); + let fix = diagnostic.fix_mut().unwrap(); + let mut edits = fix.edits().to_vec(); + for diag in diagnostics { + edits.extend_from_slice(diag.fix().unwrap().edits()); + } + *fix = Fix::unsafe_edits(edits.remove(0), edits); + + insta::assert_snapshot!(env.render(&diagnostic), @r" + error[unused-import]: `os` imported but unused + --> notebook.ipynb:cell 1:2:8 + | + 1 | # cell 1 + 2 | import os + | ^^ + | + help: Remove unused import: `os` + + ℹ Unsafe fix + ::: cell 1 + 1 1 | # cell 1 + 2 |-import os + ::: cell 2 + 1 1 | # cell 2 + 2 |-import math + 3 2 | + 4 3 | print('hello world') + ::: cell 3 + 1 1 | # cell 3 + 2 2 | def foo(): + 3 3 | print() + 4 |- x = 1 + 5 4 | + "); + } + /// Carriage return (`\r`) is a valid line-ending in Python, so we should normalize this to a /// line feed (`\n`) for rendering. Otherwise we report a single long line for this case. #[test] diff --git a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__import_sorting.snap b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__import_sorting.snap index d4cb049469..71e562fd01 100644 --- a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__import_sorting.snap +++ b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__import_sorting.snap @@ -12,14 +12,12 @@ I001 [*] Import block is un-sorted or un-formatted help: Organize imports ℹ Safe fix + ::: cell 1 1 |+import math 2 |+import random 1 3 | from pathlib import Path 2 |-import random 3 |-import math -4 4 | from typing import Any -5 5 | import collections -6 6 | # Newline should be added here I001 [*] Import block is un-sorted or un-formatted --> isort.ipynb:cell 2:1:1 @@ -33,17 +31,15 @@ I001 [*] Import block is un-sorted or un-formatted help: Organize imports ℹ Safe fix -1 1 | from pathlib import Path -2 2 | import random -3 3 | import math - 4 |+import collections -4 5 | from typing import Any -5 |-import collections - 6 |+ - 7 |+ -6 8 | # Newline should be added here -7 9 | def foo(): -8 10 | pass + ::: cell 2 + 1 |+import collections +1 2 | from typing import Any +2 |-import collections + 3 |+ + 4 |+ +3 5 | # Newline should be added here +4 6 | def foo(): +5 7 | pass I001 [*] Import block is un-sorted or un-formatted --> isort.ipynb:cell 3:1:1 @@ -57,15 +53,13 @@ I001 [*] Import block is un-sorted or un-formatted help: Organize imports ℹ Safe fix -6 6 | # Newline should be added here -7 7 | def foo(): -8 8 | pass - 9 |+import sys -9 10 | from pathlib import Path -10 |-import sys -11 11 | -12 12 | %matplotlib \ -13 13 | --inline + ::: cell 3 + 1 |+import sys +1 2 | from pathlib import Path +2 |-import sys +3 3 | +4 4 | %matplotlib \ +5 5 | --inline I001 [*] Import block is un-sorted or un-formatted --> isort.ipynb:cell 3:7:1 @@ -79,9 +73,10 @@ I001 [*] Import block is un-sorted or un-formatted help: Organize imports ℹ Safe fix -12 12 | %matplotlib \ -13 13 | --inline -14 14 | - 15 |+import abc -15 16 | import math -16 |-import abc + ::: cell 3 +4 4 | %matplotlib \ +5 5 | --inline +6 6 | + 7 |+import abc +7 8 | import math +8 |-import abc diff --git a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__ipy_escape_command.snap b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__ipy_escape_command.snap index d373e3d036..95e166dfc5 100644 --- a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__ipy_escape_command.snap +++ b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__ipy_escape_command.snap @@ -14,13 +14,13 @@ F401 [*] `os` imported but unused help: Remove unused import: `os` ℹ Safe fix + ::: cell 1 2 2 | 3 3 | %matplotlib inline 4 4 | 5 |-import os 6 5 | 7 6 | _ = math.pi -8 7 | %%timeit F401 [*] `sys` imported but unused --> ipy_escape_command.ipynb:cell 2:2:8 @@ -32,7 +32,6 @@ F401 [*] `sys` imported but unused help: Remove unused import: `sys` ℹ Safe fix -6 6 | -7 7 | _ = math.pi -8 8 | %%timeit -9 |-import sys + ::: cell 2 +1 1 | %%timeit +2 |-import sys diff --git a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__unused_variable.snap b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__unused_variable.snap index 3e51ee839d..e79da46867 100644 --- a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__unused_variable.snap +++ b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__unused_variable.snap @@ -12,12 +12,11 @@ F841 [*] Local variable `foo1` is assigned to but never used help: Remove assignment to unused variable `foo1` ℹ Unsafe fix + ::: cell 1 1 1 | def f(): 2 |- foo1 = %matplotlib --list 2 |+ %matplotlib --list 3 3 | foo2: list[str] = %matplotlib --list -4 4 | def f(): -5 5 | bar1 = !pwd F841 [*] Local variable `foo2` is assigned to but never used --> unused_variable.ipynb:cell 1:3:5 @@ -30,13 +29,11 @@ F841 [*] Local variable `foo2` is assigned to but never used help: Remove assignment to unused variable `foo2` ℹ Unsafe fix + ::: cell 1 1 1 | def f(): 2 2 | foo1 = %matplotlib --list 3 |- foo2: list[str] = %matplotlib --list 3 |+ %matplotlib --list -4 4 | def f(): -5 5 | bar1 = !pwd -6 6 | bar2: str = !pwd F841 [*] Local variable `bar1` is assigned to but never used --> unused_variable.ipynb:cell 2:2:5 @@ -49,12 +46,11 @@ F841 [*] Local variable `bar1` is assigned to but never used help: Remove assignment to unused variable `bar1` ℹ Unsafe fix -2 2 | foo1 = %matplotlib --list -3 3 | foo2: list[str] = %matplotlib --list -4 4 | def f(): -5 |- bar1 = !pwd - 5 |+ !pwd -6 6 | bar2: str = !pwd + ::: cell 2 +1 1 | def f(): +2 |- bar1 = !pwd + 2 |+ !pwd +3 3 | bar2: str = !pwd F841 [*] Local variable `bar2` is assigned to but never used --> unused_variable.ipynb:cell 2:3:5 @@ -67,8 +63,8 @@ F841 [*] Local variable `bar2` is assigned to but never used help: Remove assignment to unused variable `bar2` ℹ Unsafe fix -3 3 | foo2: list[str] = %matplotlib --list -4 4 | def f(): -5 5 | bar1 = !pwd -6 |- bar2: str = !pwd - 6 |+ !pwd + ::: cell 2 +1 1 | def f(): +2 2 | bar1 = !pwd +3 |- bar2: str = !pwd + 3 |+ !pwd diff --git a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__vscode_language_id.snap b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__vscode_language_id.snap index 8bdc926bd5..72e92de07d 100644 --- a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__vscode_language_id.snap +++ b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__vscode_language_id.snap @@ -12,6 +12,7 @@ F401 [*] `os` imported but unused help: Remove unused import: `os` ℹ Safe fix + ::: cell 3 1 |-import os 2 1 | 3 2 | print("hello world") diff --git a/crates/ruff_notebook/src/index.rs b/crates/ruff_notebook/src/index.rs index 35e4e07fcb..eff605aa6d 100644 --- a/crates/ruff_notebook/src/index.rs +++ b/crates/ruff_notebook/src/index.rs @@ -33,6 +33,14 @@ impl NotebookIndex { self.row_to_row_in_cell.get(row.to_zero_indexed()).copied() } + /// Returns an iterator over the row:cell-number pairs (both 1-based). + pub fn iter(&self) -> impl Iterator { + self.row_to_cell + .iter() + .enumerate() + .map(|(row, cell)| (OneIndexed::from_zero_indexed(row), *cell)) + } + /// Translates the given [`LineColumn`] based on the indexing table. /// /// This will translate the row/column in the concatenated source code