Consider Jupyter index for code frames (--show-source) (#5402)

## Summary

Consider Jupyter index for code frames (`--show-source`).

This solves two problems as mentioned in the linked issue:

> Omit any contents from adjoining cells

If the Jupyter index is present, we'll use that to check if the
surrounding
lines belong to the same cell as the content line. If not, we'll skip
that line
until we either reach the one which does or we reach the content line.

> code frame line number

If the Jupyter index is present, we'll use that to get the actual start
line in
corresponding to the computed start index.

## Test Plan

`cargo run --bin ruff -- check --no-cache --isolated --select=ALL --show-source /path/to/notebook.ipynb`

fixes: #5395

crates/ruff/src/message/grouped.rs

@@ -149,7 +149,14 @@ impl Display for DisplayGroupedMessage<'_> {
         if self.show_source {
             use std::fmt::Write;
             let mut padded = PadAdapter::new(f);
-            writeln!(padded, "{}", MessageCodeFrame { message })?;
+            writeln!(
+                padded,
+                "{}",
+                MessageCodeFrame {
+                    message,
+                    jupyter_index: self.jupyter_index
+                }
+            )?;
         }
 
         Ok(())

crates/ruff/src/message/text.rs

@@ -11,7 +11,7 @@ use ruff_text_size::{TextRange, TextSize};
 use ruff_python_ast::source_code::{OneIndexed, SourceLocation};
 
 use crate::fs::relativize_path;
-use crate::jupyter::Notebook;
+use crate::jupyter::{JupyterIndex, Notebook};
 use crate::line_width::{LineWidth, TabSize};
 use crate::message::diff::Diff;
 use crate::message::{Emitter, EmitterContext, Message};
@@ -72,13 +72,13 @@ impl Emitter for TextEmitter {
             )?;
 
             let start_location = message.compute_start_location();
-
-            // Check if we're working on a jupyter notebook and translate positions with cell accordingly
-            let diagnostic_location = if let Some(jupyter_index) = context
+            let jupyter_index = context
                 .source_kind(message.filename())
                 .and_then(SourceKind::notebook)
-                .map(Notebook::index)
-            {
+                .map(Notebook::index);
+
+            // Check if we're working on a jupyter notebook and translate positions with cell accordingly
+            let diagnostic_location = if let Some(jupyter_index) = jupyter_index {
                 write!(
                     writer,
                     "cell {cell}{sep}",
@@ -114,7 +114,14 @@ impl Emitter for TextEmitter {
             )?;
 
             if self.flags.contains(EmitterFlags::SHOW_SOURCE) {
-                writeln!(writer, "{}", MessageCodeFrame { message })?;
+                writeln!(
+                    writer,
+                    "{}",
+                    MessageCodeFrame {
+                        message,
+                        jupyter_index
+                    }
+                )?;
             }
 
             if self.flags.contains(EmitterFlags::SHOW_FIX_DIFF) {
@@ -158,6 +165,7 @@ impl Display for RuleCodeAndBody<'_> {
 
 pub(super) struct MessageCodeFrame<'a> {
     pub(crate) message: &'a Message,
+    pub(crate) jupyter_index: Option<&'a JupyterIndex>,
 }
 
 impl Display for MessageCodeFrame<'_> {
@@ -182,6 +190,20 @@ impl Display for MessageCodeFrame<'_> {
         let content_start_index = source_code.line_index(range.start());
         let mut start_index = content_start_index.saturating_sub(2);
 
+        // If we're working on a jupyter notebook, skip the lines which are
+        // outside of the cell containing the diagnostic.
+        if let Some(jupyter_index) = self.jupyter_index {
+            let content_start_cell = jupyter_index
+                .cell(content_start_index.get())
+                .unwrap_or_default();
+            while start_index < content_start_index {
+                if jupyter_index.cell(start_index.get()).unwrap_or_default() == content_start_cell {
+                    break;
+                }
+                start_index = start_index.saturating_add(1);
+            }
+        }
+
         // Trim leading empty lines.
         while start_index < content_start_index {
             if !source_code.line_text(start_index).trim().is_empty() {
@@ -195,7 +217,21 @@ impl Display for MessageCodeFrame<'_> {
             .saturating_add(2)
             .min(OneIndexed::from_zero_indexed(source_code.line_count()));
 
-        // Trim trailing empty lines
+        // If we're working on a jupyter notebook, skip the lines which are
+        // outside of the cell containing the diagnostic.
+        if let Some(jupyter_index) = self.jupyter_index {
+            let content_end_cell = jupyter_index
+                .cell(content_end_index.get())
+                .unwrap_or_default();
+            while end_index < content_end_index {
+                if jupyter_index.cell(end_index.get()).unwrap_or_default() == content_end_cell {
+                    break;
+                }
+                end_index = end_index.saturating_sub(1);
+            }
+        }
+
+        // Trim trailing empty lines.
         while end_index > content_end_index {
             if !source_code.line_text(end_index).trim().is_empty() {
                 break;
@@ -224,7 +260,14 @@ impl Display for MessageCodeFrame<'_> {
             title: None,
             slices: vec![Slice {
                 source: &source.text,
-                line_start: start_index.get(),
+                line_start: self.jupyter_index.map_or_else(
+                    || start_index.get(),
+                    |jupyter_index| {
+                        jupyter_index
+                            .cell_row(start_index.get())
+                            .unwrap_or_default() as usize
+                    },
+                ),
                 annotations: vec![SourceAnnotation {
                     label: &label,
                     annotation_type: AnnotationType::Error,