[ty] Remap Jupyter notebook cell indices in ruff_db (#19698)

## Summary

This PR remaps ranges in Jupyter notebooks from simple `row:column`
indices in the concatenated source code to `cell:row:col` to match
Ruff's output. This is probably not a likely change to land upstream in
`annotate-snippets`, but I didn't see a good way around it.

The remapping logic is taken nearly verbatim from here:


cd6bf1457d/crates/ruff_linter/src/message/text.rs (L212-L222)


## Test Plan

New `full` rendering test for a notebook

I was mainly focused on Ruff, but in local tests this also works for ty:

```
error[invalid-assignment]: Object of type `Literal[1]` is not assignable to `str`
 --> Untitled.ipynb:cell 1:3:1
  |
1 | import math
2 |
3 | x: str = 1
  | ^
  |
info: rule `invalid-assignment` is enabled by default

error[invalid-assignment]: Object of type `Literal[1]` is not assignable to `str`
 --> Untitled.ipynb:cell 2:3:1
  |
1 | import math
2 |
3 | x: str = 1
  | ^
  |
info: rule `invalid-assignment` is enabled by default
```

This isn't a duplicate diagnostic, just an unimaginative example:

```py
# cell 1
import math

x: str = 1
# cell 2
import math

x: str = 1
```

crates/ruff_annotate_snippets/src/renderer/display_list.rs

@@ -263,7 +263,11 @@ impl DisplaySet<'_> {
             if annotation.is_fixable {
                 buffer.append(line_offset, "[", stylesheet.none);
                 buffer.append(line_offset, "*", stylesheet.help);
-                buffer.append(line_offset, "] ", stylesheet.none);
+                buffer.append(line_offset, "]", stylesheet.none);
+                // In the hide-severity case, we need a space instead of the colon and space below.
+                if hide_severity {
+                    buffer.append(line_offset, " ", stylesheet.none);
+                }
             }
 
             if !is_annotation_empty(annotation) {
@@ -298,11 +302,15 @@ impl DisplaySet<'_> {
                 let lineno_color = stylesheet.line_no();
                 buffer.puts(line_offset, lineno_width, header_sigil, *lineno_color);
                 buffer.puts(line_offset, lineno_width + 4, path, stylesheet.none);
-                if let Some((col, row)) = pos {
+                if let Some(Position { row, col, cell }) = pos {
+                    if let Some(cell) = cell {
                         buffer.append(line_offset, ":", stylesheet.none);
-                    buffer.append(line_offset, col.to_string().as_str(), stylesheet.none);
+                        buffer.append(line_offset, &format!("cell {cell}"), stylesheet.none);
+                    }
                     buffer.append(line_offset, ":", stylesheet.none);
                     buffer.append(line_offset, row.to_string().as_str(), stylesheet.none);
+                    buffer.append(line_offset, ":", stylesheet.none);
+                    buffer.append(line_offset, col.to_string().as_str(), stylesheet.none);
                 }
                 Ok(())
             }
@@ -883,6 +891,13 @@ impl DisplaySourceAnnotation<'_> {
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub(crate) struct Position {
+    row: usize,
+    col: usize,
+    cell: Option<usize>,
+}
+
 /// Raw line - a line which does not have the `lineno` part and is not considered
 /// a part of the snippet.
 #[derive(Debug, PartialEq)]
@@ -891,7 +906,7 @@ pub(crate) enum DisplayRawLine<'a> {
     /// slice in the project structure.
     Origin {
         path: &'a str,
-        pos: Option<(usize, usize)>,
+        pos: Option<Position>,
         header_type: DisplayHeaderType,
     },
 
@@ -1191,13 +1206,15 @@ fn format_snippet<'m>(
             "Non-empty file-level snippet that won't be rendered: {:?}",
             snippet.source
         );
-        let header = format_header(origin, main_range, &[], is_first);
+        let header = format_header(origin, main_range, &[], is_first, snippet.cell_index);
         return DisplaySet {
             display_lines: header.map_or_else(Vec::new, |header| vec![header]),
             margin: Margin::new(0, 0, 0, 0, term_width, 0),
         };
     }
 
+    let cell_index = snippet.cell_index;
+
     let mut body = format_body(
         snippet,
         need_empty_header,
@@ -1206,7 +1223,13 @@ fn format_snippet<'m>(
         anonymized_line_numbers,
         cut_indicator,
     );
-    let header = format_header(origin, main_range, &body.display_lines, is_first);
+    let header = format_header(
+        origin,
+        main_range,
+        &body.display_lines,
+        is_first,
+        cell_index,
+    );
 
     if let Some(header) = header {
         body.display_lines.insert(0, header);
@@ -1226,6 +1249,7 @@ fn format_header<'a>(
     main_range: Option<usize>,
     body: &[DisplayLine<'_>],
     is_first: bool,
+    cell_index: Option<usize>,
 ) -> Option<DisplayLine<'a>> {
     let display_header = if is_first {
         DisplayHeaderType::Initial
@@ -1262,7 +1286,11 @@ fn format_header<'a>(
 
         return Some(DisplayLine::Raw(DisplayRawLine::Origin {
             path,
-            pos: Some((line_offset, col)),
+            pos: Some(Position {
+                row: line_offset,
+                col,
+                cell: cell_index,
+            }),
             header_type: display_header,
         }));
     }

crates/ruff_annotate_snippets/src/snippet.rs

@@ -75,6 +75,10 @@ pub struct Snippet<'a> {
     pub(crate) annotations: Vec<Annotation<'a>>,
 
     pub(crate) fold: bool,
+
+    /// The optional cell index in a Jupyter notebook, used for reporting source locations along
+    /// with the ranges on `annotations`.
+    pub(crate) cell_index: Option<usize>,
 }
 
 impl<'a> Snippet<'a> {
@@ -85,6 +89,7 @@ impl<'a> Snippet<'a> {
             source,
             annotations: vec![],
             fold: false,
+            cell_index: None,
         }
     }
 
@@ -113,6 +118,12 @@ impl<'a> Snippet<'a> {
         self.fold = fold;
         self
     }
+
+    /// Attach a Jupyter notebook cell index.
+    pub fn cell_index(mut self, index: Option<usize>) -> Self {
+        self.cell_index = index;
+        self
+    }
 }
 
 /// An annotation for a [`Snippet`].

crates/ruff_db/src/diagnostic/render.rs

@@ -244,7 +244,7 @@ impl<'a> ResolvedDiagnostic<'a> {
             .filter_map(|ann| {
                 let path = ann.span.file.path(resolver);
                 let diagnostic_source = ann.span.file.diagnostic_source(resolver);
-                ResolvedAnnotation::new(path, &diagnostic_source, ann)
+                ResolvedAnnotation::new(path, &diagnostic_source, ann, resolver)
             })
             .collect();
 
@@ -291,7 +291,7 @@ impl<'a> ResolvedDiagnostic<'a> {
             .filter_map(|ann| {
                 let path = ann.span.file.path(resolver);
                 let diagnostic_source = ann.span.file.diagnostic_source(resolver);
-                ResolvedAnnotation::new(path, &diagnostic_source, ann)
+                ResolvedAnnotation::new(path, &diagnostic_source, ann, resolver)
             })
             .collect();
         ResolvedDiagnostic {
@@ -330,20 +330,49 @@ impl<'a> ResolvedDiagnostic<'a> {
                     &prev.diagnostic_source.as_source_code(),
                     context,
                     prev.line_end,
+                    prev.notebook_index.as_ref(),
                 )
                 .get();
                 let this_context_begins = context_before(
                     &ann.diagnostic_source.as_source_code(),
                     context,
                     ann.line_start,
+                    ann.notebook_index.as_ref(),
                 )
                 .get();
+
+                // For notebooks, check whether the end of the
+                // previous annotation and the start of the current
+                // annotation are in different cells.
+                let prev_cell_index = prev.notebook_index.as_ref().map(|notebook_index| {
+                    let prev_end = prev
+                        .diagnostic_source
+                        .as_source_code()
+                        .line_column(prev.range.end());
+                    notebook_index.cell(prev_end.line).unwrap_or_default().get()
+                });
+                let this_cell_index = ann.notebook_index.as_ref().map(|notebook_index| {
+                    let this_start = ann
+                        .diagnostic_source
+                        .as_source_code()
+                        .line_column(ann.range.start());
+                    notebook_index
+                        .cell(this_start.line)
+                        .unwrap_or_default()
+                        .get()
+                });
+                let in_different_cells = prev_cell_index != this_cell_index;
+
                 // The boundary case here is when `prev_context_ends`
                 // is exactly one less than `this_context_begins`. In
                 // that case, the context windows are adjacent and we
                 // should fall through below to add this annotation to
                 // the existing snippet.
-                if this_context_begins.saturating_sub(prev_context_ends) > 1 {
+                //
+                // For notebooks, also check that the context windows
+                // are in the same cell. Windows from different cells
+                // should never be considered adjacent.
+                if in_different_cells || this_context_begins.saturating_sub(prev_context_ends) > 1 {
                     snippet_by_path
                         .entry(path)
                         .or_default()
@@ -388,6 +417,7 @@ struct ResolvedAnnotation<'a> {
     message: Option<&'a str>,
     is_primary: bool,
     is_file_level: bool,
+    notebook_index: Option<NotebookIndex>,
 }
 
 impl<'a> ResolvedAnnotation<'a> {
@@ -400,6 +430,7 @@ impl<'a> ResolvedAnnotation<'a> {
         path: &'a str,
         diagnostic_source: &DiagnosticSource,
         ann: &'a Annotation,
+        resolver: &'a dyn FileResolver,
     ) -> Option<ResolvedAnnotation<'a>> {
         let source = diagnostic_source.as_source_code();
         let (range, line_start, line_end) = match (ann.span.range(), ann.message.is_some()) {
@@ -434,6 +465,7 @@ impl<'a> ResolvedAnnotation<'a> {
             message: ann.get_message(),
             is_primary: ann.is_primary,
             is_file_level: ann.is_file_level,
+            notebook_index: resolver.notebook_index(&ann.span.file),
         })
     }
 }
@@ -566,17 +598,27 @@ struct RenderableSnippet<'r> {
     /// Whether this snippet contains at least one primary
     /// annotation.
     has_primary: bool,
+    /// The cell index in a Jupyter notebook, if this snippet refers to a notebook.
+    ///
+    /// This is used for rendering annotations with offsets like `cell 1:2:3` instead of simple row
+    /// and column numbers.
+    cell_index: Option<usize>,
 }
 
 impl<'r> RenderableSnippet<'r> {
     /// Creates a new snippet with one or more annotations that is ready to be
-    /// renderer.
+    /// rendered.
     ///
     /// The first line of the snippet is the smallest line number on which one
     /// of the annotations begins, minus the context window size. The last line
     /// is the largest line number on which one of the annotations ends, plus
     /// the context window size.
     ///
+    /// For Jupyter notebooks, the context window may also be truncated at cell
+    /// boundaries. If multiple annotations are present, and they point to
+    /// different cells, these will have already been split into separate
+    /// snippets by `ResolvedDiagnostic::to_renderable`.
+    ///
     /// Callers should guarantee that the `input` on every `ResolvedAnnotation`
     /// given is identical.
     ///
@@ -593,19 +635,19 @@ impl<'r> RenderableSnippet<'r> {
             "creating a renderable snippet requires a non-zero number of annotations",
         );
         let diagnostic_source = &anns[0].diagnostic_source;
+        let notebook_index = anns[0].notebook_index.as_ref();
         let source = diagnostic_source.as_source_code();
         let has_primary = anns.iter().any(|ann| ann.is_primary);
 
-        let line_start = context_before(
-            &source,
-            context,
-            anns.iter().map(|ann| ann.line_start).min().unwrap(),
-        );
-        let line_end = context_after(
-            &source,
-            context,
-            anns.iter().map(|ann| ann.line_end).max().unwrap(),
-        );
+        let content_start_index = anns.iter().map(|ann| ann.line_start).min().unwrap();
+        let line_start = context_before(&source, context, content_start_index, notebook_index);
+
+        let start = source.line_column(anns[0].range.start());
+        let cell_index = notebook_index
+            .map(|notebook_index| notebook_index.cell(start.line).unwrap_or_default().get());
+
+        let content_end_index = anns.iter().map(|ann| ann.line_end).max().unwrap();
+        let line_end = context_after(&source, context, content_end_index, notebook_index);
 
         let snippet_start = source.line_start(line_start);
         let snippet_end = source.line_end(line_end);
@@ -623,11 +665,18 @@ impl<'r> RenderableSnippet<'r> {
             annotations,
         } = replace_unprintable(snippet, annotations).fix_up_empty_spans_after_line_terminator();
 
+        let line_start = notebook_index.map_or(line_start, |notebook_index| {
+            notebook_index
+                .cell_row(line_start)
+                .unwrap_or(OneIndexed::MIN)
+        });
+
         RenderableSnippet {
             snippet,
             line_start,
             annotations,
             has_primary,
+            cell_index,
         }
     }
 
@@ -641,6 +690,7 @@ impl<'r> RenderableSnippet<'r> {
                     .iter()
                     .map(RenderableAnnotation::to_annotate),
             )
+            .cell_index(self.cell_index)
     }
 }
 
@@ -827,7 +877,15 @@ pub struct Input {
 ///
 /// The line number returned is guaranteed to be less than
 /// or equal to `start`.
-fn context_before(source: &SourceCode<'_, '_>, len: usize, start: OneIndexed) -> OneIndexed {
+///
+/// In Jupyter notebooks, lines outside the cell containing
+/// `start` will be omitted.
+fn context_before(
+    source: &SourceCode<'_, '_>,
+    len: usize,
+    start: OneIndexed,
+    notebook_index: Option<&NotebookIndex>,
+) -> OneIndexed {
     let mut line = start.saturating_sub(len);
     // Trim leading empty lines.
     while line < start {
@@ -836,6 +894,17 @@ fn context_before(source: &SourceCode<'_, '_>, len: usize, start: OneIndexed) ->
         }
         line = line.saturating_add(1);
     }
+
+    if let Some(index) = notebook_index {
+        let content_start_cell = index.cell(start).unwrap_or(OneIndexed::MIN);
+        while line < start {
+            if index.cell(line).unwrap_or(OneIndexed::MIN) == content_start_cell {
+                break;
+            }
+            line = line.saturating_add(1);
+        }
+    }
+
     line
 }
 
@@ -845,7 +914,15 @@ fn context_before(source: &SourceCode<'_, '_>, len: usize, start: OneIndexed) ->
 /// The line number returned is guaranteed to be greater
 /// than or equal to `start` and no greater than the
 /// number of lines in `source`.
-fn context_after(source: &SourceCode<'_, '_>, len: usize, start: OneIndexed) -> OneIndexed {
+///
+/// In Jupyter notebooks, lines outside the cell containing
+/// `start` will be omitted.
+fn context_after(
+    source: &SourceCode<'_, '_>,
+    len: usize,
+    start: OneIndexed,
+    notebook_index: Option<&NotebookIndex>,
+) -> OneIndexed {
     let max_lines = OneIndexed::from_zero_indexed(source.line_count());
     let mut line = start.saturating_add(len).min(max_lines);
     // Trim trailing empty lines.
@@ -855,6 +932,17 @@ fn context_after(source: &SourceCode<'_, '_>, len: usize, start: OneIndexed) ->
         }
         line = line.saturating_sub(1);
     }
+
+    if let Some(index) = notebook_index {
+        let content_end_cell = index.cell(start).unwrap_or(OneIndexed::MIN);
+        while line > start {
+            if index.cell(line).unwrap_or(OneIndexed::MIN) == content_end_cell {
+                break;
+            }
+            line = line.saturating_sub(1);
+        }
+    }
+
     line
 }
 
@@ -2698,7 +2786,7 @@ watermelon
         ///
         /// See the docs on `TestEnvironment::span` for the meaning of
         /// `path`, `line_offset_start` and `line_offset_end`.
-        fn secondary(
+        pub(super) fn secondary(
             mut self,
             path: &str,
             line_offset_start: &str,
@@ -2734,7 +2822,7 @@ watermelon
         }
 
         /// Adds a "help" sub-diagnostic with the given message.
-        fn help(mut self, message: impl IntoDiagnosticMessage) -> DiagnosticBuilder<'e> {
+        pub(super) fn help(mut self, message: impl IntoDiagnosticMessage) -> DiagnosticBuilder<'e> {
             self.diag.help(message);
             self
         }
@@ -2905,7 +2993,8 @@ if call(foo
         (env, diagnostics)
     }
 
-    /// Create Ruff-style diagnostics for testing the various output formats for a notebook.
+    /// A Jupyter notebook for testing diagnostics.
+    ///
     ///
     /// The concatenated cells look like this:
     ///
@@ -2925,17 +3014,7 @@ if call(foo
     /// The first diagnostic is on the unused `os` import with location cell 1, row 2, column 8
     /// (`cell 1:2:8`). The second diagnostic is the unused `math` import at `cell 2:2:8`, and the
     /// third diagnostic is an unfixable unused variable at `cell 3:4:5`.
-    #[allow(
-        dead_code,
-        reason = "This is currently only used for JSON but will be needed soon for other formats"
-    )]
-    pub(crate) fn create_notebook_diagnostics(
-        format: DiagnosticFormat,
-    ) -> (TestEnvironment, Vec<Diagnostic>) {
-        let mut env = TestEnvironment::new();
-        env.add(
-            "notebook.ipynb",
-            r##"
+    pub(super) static NOTEBOOK: &str = r##"
         {
  "cells": [
   {
@@ -2974,8 +3053,14 @@ if call(foo
  "nbformat": 4,
  "nbformat_minor": 5
 }
-"##,
-        );
+"##;
+
+    /// Create Ruff-style diagnostics for testing the various output formats for a notebook.
+    pub(crate) fn create_notebook_diagnostics(
+        format: DiagnosticFormat,
+    ) -> (TestEnvironment, Vec<Diagnostic>) {
+        let mut env = TestEnvironment::new();
+        env.add("notebook.ipynb", NOTEBOOK);
         env.format(format);
 
         let diagnostics = vec![

crates/ruff_db/src/diagnostic/render/full.rs

@@ -5,7 +5,10 @@ mod tests {
 
     use crate::diagnostic::{
         Annotation, DiagnosticFormat, Severity,
-        render::tests::{TestEnvironment, create_diagnostics, create_syntax_error_diagnostics},
+        render::tests::{
+            NOTEBOOK, TestEnvironment, create_diagnostics, create_notebook_diagnostics,
+            create_syntax_error_diagnostics,
+        },
     };
 
     #[test]
@@ -285,4 +288,116 @@ print()
         --> example.py:1:1
         ");
     }
+
+    /// Check that ranges in notebooks are remapped relative to the cells.
+    #[test]
+    fn notebook_output() {
+        let (env, diagnostics) = create_notebook_diagnostics(DiagnosticFormat::Full);
+        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`
+
+        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`
+
+        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`
+        ");
+    }
+
+    /// Check notebook handling for multiple annotations in a single diagnostic that span cells.
+    #[test]
+    fn notebook_output_multiple_annotations() {
+        let mut env = TestEnvironment::new();
+        env.add("notebook.ipynb", NOTEBOOK);
+
+        let diagnostics = vec![
+            // adjacent context windows
+            env.builder("unused-import", Severity::Error, "`os` imported but unused")
+                .primary("notebook.ipynb", "2:7", "2:9", "")
+                .secondary("notebook.ipynb", "4:7", "4:11", "second cell")
+                .help("Remove unused import: `os`")
+                .build(),
+            // non-adjacent context windows
+            env.builder("unused-import", Severity::Error, "`os` imported but unused")
+                .primary("notebook.ipynb", "2:7", "2:9", "")
+                .secondary("notebook.ipynb", "10:4", "10:5", "second cell")
+                .help("Remove unused import: `os`")
+                .build(),
+            // adjacent context windows in the same cell
+            env.err()
+                .primary("notebook.ipynb", "4:7", "4:11", "second cell")
+                .secondary("notebook.ipynb", "6:0", "6:5", "print statement")
+                .help("Remove `print` statement")
+                .build(),
+        ];
+
+        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
+          |        ^^
+          |
+         ::: notebook.ipynb:cell 2:2:8
+          |
+        1 | # cell 2
+        2 | import math
+          |        ---- second cell
+        3 |
+        4 | print('hello world')
+          |
+        help: Remove unused import: `os`
+
+        error[unused-import]: `os` imported but unused
+         --> notebook.ipynb:cell 1:2:8
+          |
+        1 | # cell 1
+        2 | import os
+          |        ^^
+          |
+         ::: notebook.ipynb:cell 3:4:5
+          |
+        2 | def foo():
+        3 |     print()
+        4 |     x = 1
+          |     - second cell
+          |
+        help: Remove unused import: `os`
+
+        error[test-diagnostic]: main diagnostic message
+         --> notebook.ipynb:cell 2:2:8
+          |
+        1 | # cell 2
+        2 | import math
+          |        ^^^^ second cell
+        3 |
+        4 | print('hello world')
+          | ----- print statement
+          |
+        help: Remove `print` statement
+        ");
+    }
 }