Add test for Notebook text output

crates/ruff_linter/src/logging.rs

@@ -177,18 +177,15 @@ impl Display for DisplayParseError<'_> {
                     f,
                     "cell {cell}{colon}",
                     cell = jupyter_index
-                        .cell(source_location.row.get())
-                        .unwrap_or_default(),
+                        .cell(source_location.row)
+                        .unwrap_or(OneIndexed::MIN),
                     colon = ":".cyan(),
                 )?;
 
                 SourceLocation {
-                    row: OneIndexed::new(
-                        jupyter_index
-                            .cell_row(source_location.row.get())
-                            .unwrap_or(1) as usize,
-                    )
-                    .unwrap(),
+                    row: jupyter_index
+                        .cell_row(source_location.row)
+                        .unwrap_or(OneIndexed::MIN),
                     column: source_location.column,
                 }
             } else {

crates/ruff_linter/src/message/diff.rs

@@ -35,6 +35,7 @@ impl<'a> Diff<'a> {
 
 impl Display for Diff<'_> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        // TODO(dhruvmanila): Add support for Notebook cells once it's user-facing
         let mut output = String::with_capacity(self.source_code.source_text().len());
         let mut last_end = TextSize::default();
 

crates/ruff_linter/src/message/grouped.rs

@@ -125,18 +125,18 @@ impl Display for DisplayGroupedMessage<'_> {
                 f,
                 "cell {cell}{sep}",
                 cell = jupyter_index
-                    .cell(start_location.row.get())
-                    .unwrap_or_default(),
+                    .cell(start_location.row)
+                    .unwrap_or(OneIndexed::MIN),
                 sep = ":".cyan()
             )?;
             (
                 jupyter_index
-                    .cell_row(start_location.row.get())
-                    .unwrap_or(1) as usize,
-                start_location.column.get(),
+                    .cell_row(start_location.row)
+                    .unwrap_or(OneIndexed::MIN),
+                start_location.column,
             )
         } else {
-            (start_location.row.get(), start_location.column.get())
+            (start_location.row, start_location.column)
         };
 
         writeln!(

crates/ruff_linter/src/message/mod.rs

@@ -150,7 +150,8 @@ mod tests {
     use rustc_hash::FxHashMap;
 
     use ruff_diagnostics::{Diagnostic, DiagnosticKind, Edit, Fix};
-    use ruff_source_file::SourceFileBuilder;
+    use ruff_notebook::NotebookIndex;
+    use ruff_source_file::{OneIndexed, SourceFileBuilder};
     use ruff_text_size::{Ranged, TextRange, TextSize};
 
     use crate::message::{Emitter, EmitterContext, Message};
@@ -221,6 +222,113 @@ def fibonacci(n):
         ]
     }
 
+    pub(super) fn create_notebook_messages() -> (Vec<Message>, FxHashMap<String, NotebookIndex>) {
+        let notebook = r"# cell 1
+import os
+# cell 2
+import math
+
+print('hello world')
+# cell 3
+def foo():
+    print()
+    x = 1
+";
+
+        let unused_import_os = Diagnostic::new(
+            DiagnosticKind {
+                name: "UnusedImport".to_string(),
+                body: "`os` imported but unused".to_string(),
+                suggestion: Some("Remove unused import: `os`".to_string()),
+            },
+            TextRange::new(TextSize::from(16), TextSize::from(18)),
+        )
+        .with_fix(Fix::safe_edit(Edit::range_deletion(TextRange::new(
+            TextSize::from(9),
+            TextSize::from(19),
+        ))));
+
+        let unused_import_math = Diagnostic::new(
+            DiagnosticKind {
+                name: "UnusedImport".to_string(),
+                body: "`math` imported but unused".to_string(),
+                suggestion: Some("Remove unused import: `math`".to_string()),
+            },
+            TextRange::new(TextSize::from(35), TextSize::from(39)),
+        )
+        .with_fix(Fix::safe_edit(Edit::range_deletion(TextRange::new(
+            TextSize::from(28),
+            TextSize::from(40),
+        ))));
+
+        let unused_variable = Diagnostic::new(
+            DiagnosticKind {
+                name: "UnusedVariable".to_string(),
+                body: "Local variable `x` is assigned to but never used".to_string(),
+                suggestion: Some("Remove assignment to unused variable `x`".to_string()),
+            },
+            TextRange::new(TextSize::from(98), TextSize::from(99)),
+        )
+        .with_fix(Fix::unsafe_edit(Edit::deletion(
+            TextSize::from(94),
+            TextSize::from(104),
+        )));
+
+        let notebook_source = SourceFileBuilder::new("notebook.ipynb", notebook).finish();
+
+        let mut notebook_indexes = FxHashMap::default();
+        notebook_indexes.insert(
+            "notebook.ipynb".to_string(),
+            NotebookIndex::new(
+                vec![
+                    OneIndexed::from_zero_indexed(0),
+                    OneIndexed::from_zero_indexed(0),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(2),
+                    OneIndexed::from_zero_indexed(2),
+                    OneIndexed::from_zero_indexed(2),
+                    OneIndexed::from_zero_indexed(2),
+                ],
+                vec![
+                    OneIndexed::from_zero_indexed(0),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(0),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(2),
+                    OneIndexed::from_zero_indexed(3),
+                    OneIndexed::from_zero_indexed(0),
+                    OneIndexed::from_zero_indexed(1),
+                    OneIndexed::from_zero_indexed(2),
+                    OneIndexed::from_zero_indexed(3),
+                ],
+            ),
+        );
+
+        let unused_import_os_start = unused_import_os.start();
+        let unused_import_math_start = unused_import_math.start();
+        let unused_variable_start = unused_variable.start();
+
+        (
+            vec![
+                Message::from_diagnostic(
+                    unused_import_os,
+                    notebook_source.clone(),
+                    unused_import_os_start,
+                ),
+                Message::from_diagnostic(
+                    unused_import_math,
+                    notebook_source.clone(),
+                    unused_import_math_start,
+                ),
+                Message::from_diagnostic(unused_variable, notebook_source, unused_variable_start),
+            ],
+            notebook_indexes,
+        )
+    }
+
     pub(super) fn capture_emitter_output(
         emitter: &mut dyn Emitter,
         messages: &[Message],
@@ -232,4 +340,16 @@ def fibonacci(n):
 
         String::from_utf8(output).expect("Output to be valid UTF-8")
     }
+
+    pub(super) fn capture_emitter_notebook_output(
+        emitter: &mut dyn Emitter,
+        messages: &[Message],
+        notebook_indexes: &FxHashMap<String, NotebookIndex>,
+    ) -> String {
+        let context = EmitterContext::new(notebook_indexes);
+        let mut output: Vec<u8> = Vec::new();
+        emitter.emit(&mut output, messages, &context).unwrap();
+
+        String::from_utf8(output).expect("Output to be valid UTF-8")
+    }
 }

crates/ruff_linter/src/message/snapshots/ruff_linter__message__text__tests__notebook_output.snap

@@ -0,0 +1,32 @@
+---
+source: crates/ruff_linter/src/message/text.rs
+expression: content
+---
+notebook.ipynb:cell 1:2:8: F401 [*] `os` imported but unused
+  |
+1 | # cell 1
+2 | import os
+  |        ^^ F401
+  |
+  = help: Remove unused import: `os`
+
+notebook.ipynb:cell 2:2:8: F401 [*] `math` imported but unused
+  |
+1 | # cell 2
+2 | import math
+  |        ^^^^ F401
+3 | 
+4 | print('hello world')
+  |
+  = help: Remove unused import: `math`
+
+notebook.ipynb:cell 3:4:5: F841 [*] Local variable `x` is assigned to but never used
+  |
+2 | def foo():
+3 |     print()
+4 |     x = 1
+  |     ^ F841
+  |
+  = help: Remove assignment to unused variable `x`
+
+

crates/ruff_linter/src/message/text.rs

@@ -87,18 +87,15 @@ impl Emitter for TextEmitter {
                     writer,
                     "cell {cell}{sep}",
                     cell = notebook_index
-                        .cell(start_location.row.get())
-                        .unwrap_or_default(),
+                        .cell(start_location.row)
+                        .unwrap_or(OneIndexed::MIN),
                     sep = ":".cyan(),
                 )?;
 
                 SourceLocation {
-                    row: OneIndexed::new(
-                        notebook_index
-                            .cell_row(start_location.row.get())
-                            .unwrap_or(1) as usize,
-                    )
-                    .unwrap(),
+                    row: notebook_index
+                        .cell_row(start_location.row)
+                        .unwrap_or(OneIndexed::MIN),
                     column: start_location.column,
                 }
             } else {
@@ -203,9 +200,9 @@ impl Display for MessageCodeFrame<'_> {
         // If we're working with a Jupyter Notebook, skip the lines which are
         // outside of the cell containing the diagnostic.
         if let Some(index) = self.notebook_index {
-            let content_start_cell = index.cell(content_start_index.get()).unwrap_or_default();
+            let content_start_cell = index.cell(content_start_index).unwrap_or(OneIndexed::MIN);
             while start_index < content_start_index {
-                if index.cell(start_index.get()).unwrap_or_default() == content_start_cell {
+                if index.cell(start_index).unwrap_or(OneIndexed::MIN) == content_start_cell {
                     break;
                 }
                 start_index = start_index.saturating_add(1);
@@ -228,9 +225,9 @@ impl Display for MessageCodeFrame<'_> {
         // If we're working with a Jupyter Notebook, skip the lines which are
         // outside of the cell containing the diagnostic.
         if let Some(index) = self.notebook_index {
-            let content_end_cell = index.cell(content_end_index.get()).unwrap_or_default();
+            let content_end_cell = index.cell(content_end_index).unwrap_or(OneIndexed::MIN);
             while end_index > content_end_index {
-                if index.cell(end_index.get()).unwrap_or_default() == content_end_cell {
+                if index.cell(end_index).unwrap_or(OneIndexed::MIN) == content_end_cell {
                     break;
                 }
                 end_index = end_index.saturating_sub(1);
@@ -270,8 +267,9 @@ impl Display for MessageCodeFrame<'_> {
                     || start_index.get(),
                     |notebook_index| {
                         notebook_index
-                            .cell_row(start_index.get())
-                            .unwrap_or_default() as usize
+                            .cell_row(start_index)
+                            .unwrap_or(OneIndexed::MIN)
+                            .get()
                     },
                 ),
                 annotations: vec![SourceAnnotation {
@@ -353,7 +351,10 @@ struct SourceCode<'a> {
 mod tests {
     use insta::assert_snapshot;
 
-    use crate::message::tests::{capture_emitter_output, create_messages};
+    use crate::message::tests::{
+        capture_emitter_notebook_output, capture_emitter_output, create_messages,
+        create_notebook_messages,
+    };
     use crate::message::TextEmitter;
     use crate::settings::types::UnsafeFixes;
 
@@ -385,4 +386,16 @@ mod tests {
 
         assert_snapshot!(content);
     }
+
+    #[test]
+    fn notebook_output() {
+        let mut emitter = TextEmitter::default()
+            .with_show_fix_status(true)
+            .with_show_source(true)
+            .with_unsafe_fixes(UnsafeFixes::Enabled);
+        let (messages, notebook_indexes) = create_notebook_messages();
+        let content = capture_emitter_notebook_output(&mut emitter, &messages, &notebook_indexes);
+
+        assert_snapshot!(content);
+    }
 }

crates/ruff_notebook/Cargo.toml

@@ -14,7 +14,7 @@ license = { workspace = true }
 
 [dependencies]
 ruff_diagnostics = { path = "../ruff_diagnostics" }
-ruff_source_file = { path = "../ruff_source_file" }
+ruff_source_file = { path = "../ruff_source_file", features = ["serde"] }
 ruff_text_size = { path = "../ruff_text_size" }
 
 anyhow = { workspace = true }

crates/ruff_notebook/src/index.rs

@@ -1,26 +1,35 @@
 use serde::{Deserialize, Serialize};
 
+use ruff_source_file::OneIndexed;
+
 /// Jupyter Notebook indexing table
 ///
 /// When we lint a jupyter notebook, we have to translate the row/column based on
 /// [`ruff_text_size::TextSize`] to jupyter notebook cell/row/column.
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct NotebookIndex {
     /// Enter a row (1-based), get back the cell (1-based)
-    pub(super) row_to_cell: Vec<u32>,
+    pub(super) row_to_cell: Vec<OneIndexed>,
     /// Enter a row (1-based), get back the row in cell (1-based)
-    pub(super) row_to_row_in_cell: Vec<u32>,
+    pub(super) row_to_row_in_cell: Vec<OneIndexed>,
 }
 
 impl NotebookIndex {
+    pub fn new(row_to_cell: Vec<OneIndexed>, row_to_row_in_cell: Vec<OneIndexed>) -> Self {
+        Self {
+            row_to_cell,
+            row_to_row_in_cell,
+        }
+    }
+
     /// Returns the cell number (1-based) for the given row (1-based).
-    pub fn cell(&self, row: usize) -> Option<u32> {
-        self.row_to_cell.get(row).copied()
+    pub fn cell(&self, row: OneIndexed) -> Option<OneIndexed> {
+        self.row_to_cell.get(row.to_zero_indexed()).copied()
     }
 
     /// Returns the row number (1-based) in the cell (1-based) for the
     /// given row (1-based).
-    pub fn cell_row(&self, row: usize) -> Option<u32> {
-        self.row_to_row_in_cell.get(row).copied()
+    pub fn cell_row(&self, row: OneIndexed) -> Option<OneIndexed> {
+        self.row_to_row_in_cell.get(row.to_zero_indexed()).copied()
     }
 }