Improve diff rendering for notebooks (#20036)

## 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`
This commit is contained in:
Brent Westbrook 2025-08-25 09:20:42 -04:00 committed by GitHub
parent f9bbee33f6
commit d0bcf56bd9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 270 additions and 114 deletions

View File

@ -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!(

View File

@ -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);

View File

@ -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<NotebookIndex>,
stylesheet: &'a DiagnosticStylesheet,
}
@ -90,12 +92,11 @@ impl<'a> Diff<'a> {
stylesheet: &'a DiagnosticStylesheet,
resolver: &'a dyn FileResolver,
) -> Option<Diff<'a>> {
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<OneIndexed>, 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]

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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")

View File

@ -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<Item = (OneIndexed, OneIndexed)> {
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