Add CellOffsets abstraction (#8814)

Refactor `Notebook::cell_offsets` to use an abstract struct for storing
the cell offsets. This will allow us to add useful methods on it.

crates/ruff_linter/src/checkers/imports.rs

@@ -3,6 +3,7 @@ use std::borrow::Cow;
 use std::path::Path;
 
 use ruff_diagnostics::Diagnostic;
+use ruff_notebook::CellOffsets;
 use ruff_python_ast::helpers::to_module_path;
 use ruff_python_ast::imports::{ImportMap, ModuleImport};
 use ruff_python_ast::statement_visitor::StatementVisitor;
@@ -17,7 +18,6 @@ use crate::registry::Rule;
 use crate::rules::isort;
 use crate::rules::isort::block::{Block, BlockBuilder};
 use crate::settings::LinterSettings;
-use crate::source_kind::SourceKind;
 
 fn extract_import_map(path: &Path, package: Option<&Path>, blocks: &[&Block]) -> Option<ImportMap> {
     let Some(package) = package else {
@@ -85,13 +85,13 @@ pub(crate) fn check_imports(
     stylist: &Stylist,
     path: &Path,
     package: Option<&Path>,
-    source_kind: &SourceKind,
     source_type: PySourceType,
+    cell_offsets: Option<&CellOffsets>,
 ) -> (Vec<Diagnostic>, Option<ImportMap>) {
     // Extract all import blocks from the AST.
     let tracker = {
         let mut tracker =
-            BlockBuilder::new(locator, directives, source_type.is_stub(), source_kind);
+            BlockBuilder::new(locator, directives, source_type.is_stub(), cell_offsets);
         tracker.visit_body(python_ast);
         tracker
     };

crates/ruff_linter/src/linter.rs

@@ -9,6 +9,7 @@ use log::error;
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::Diagnostic;
+use ruff_notebook::Notebook;
 use ruff_python_ast::imports::ImportMap;
 use ruff_python_ast::PySourceType;
 use ruff_python_codegen::Stylist;
@@ -149,6 +150,7 @@ pub fn check_path(
             source_type.is_ipynb(),
         ) {
             Ok(python_ast) => {
+                let cell_offsets = source_kind.as_ipy_notebook().map(Notebook::cell_offsets);
                 if use_ast {
                     diagnostics.extend(check_ast(
                         &python_ast,
@@ -173,8 +175,8 @@ pub fn check_path(
                         stylist,
                         path,
                         package,
-                        source_kind,
                         source_type,
+                        cell_offsets,
                     );
                     imports = module_imports;
                     diagnostics.extend(import_diagnostics);

crates/ruff_linter/src/rules/isort/block.rs

@@ -1,7 +1,7 @@
 use std::iter::Peekable;
 use std::slice;
 
-use ruff_notebook::Notebook;
+use ruff_notebook::CellOffsets;
 use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_ast::{self as ast, ElifElseClause, ExceptHandler, MatchCase, Stmt};
 use ruff_source_file::Locator;
@@ -9,7 +9,6 @@ use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::directives::IsortDirectives;
 use crate::rules::isort::helpers;
-use crate::source_kind::SourceKind;
 
 /// A block of imports within a Python module.
 #[derive(Debug, Default)]
@@ -43,7 +42,7 @@ impl<'a> BlockBuilder<'a> {
         locator: &'a Locator<'a>,
         directives: &'a IsortDirectives,
         is_stub: bool,
-        source_kind: &'a SourceKind,
+        cell_offsets: Option<&'a CellOffsets>,
     ) -> Self {
         Self {
             locator,
@@ -52,10 +51,7 @@ impl<'a> BlockBuilder<'a> {
             splits: directives.splits.iter().peekable(),
             exclusions: &directives.exclusions,
             nested: false,
-            cell_offsets: source_kind
-                .as_ipy_notebook()
-                .map(Notebook::cell_offsets)
-                .map(|offsets| offsets.iter().peekable()),
+            cell_offsets: cell_offsets.map(|offsets| offsets.iter().peekable()),
         }
     }
 

crates/ruff_notebook/src/cell.rs

@@ -1,4 +1,7 @@
 use std::fmt;
+use std::ops::{Deref, DerefMut};
+
+use ruff_text_size::TextSize;
 
 use crate::schema::{Cell, SourceValue};
 
@@ -168,3 +171,34 @@ impl Cell {
         lines.any(|line| line.trim_start().starts_with("%%"))
     }
 }
+
+/// Cell offsets are used to keep track of the start and end offsets of each
+/// cell in the concatenated source code.
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct CellOffsets(Vec<TextSize>);
+
+impl CellOffsets {
+    /// Create a new [`CellOffsets`] with the given capacity.
+    pub(crate) fn with_capacity(capacity: usize) -> Self {
+        Self(Vec::with_capacity(capacity))
+    }
+
+    /// Push a new offset to the end of the [`CellOffsets`].
+    pub(crate) fn push(&mut self, offset: TextSize) {
+        self.0.push(offset);
+    }
+}
+
+impl Deref for CellOffsets {
+    type Target = [TextSize];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for CellOffsets {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

crates/ruff_notebook/src/lib.rs

@@ -1,5 +1,6 @@
 //! Utils for reading and writing jupyter notebooks
 
+pub use cell::*;
 pub use index::*;
 pub use notebook::*;
 pub use schema::*;

crates/ruff_notebook/src/notebook.rs

@@ -15,6 +15,7 @@ use ruff_diagnostics::{SourceMap, SourceMarker};
 use ruff_source_file::{NewlineWithTrailingNewline, OneIndexed, UniversalNewlineIterator};
 use ruff_text_size::TextSize;
 
+use crate::cell::CellOffsets;
 use crate::index::NotebookIndex;
 use crate::schema::{Cell, RawNotebook, SortAlphabetically, SourceValue};
 
@@ -65,7 +66,7 @@ pub struct Notebook {
     raw: RawNotebook,
     /// The offsets of each cell in the concatenated source code. This includes
     /// the first and last character offsets as well.
-    cell_offsets: Vec<TextSize>,
+    cell_offsets: CellOffsets,
     /// The cell index of all valid code cells in the notebook.
     valid_code_cells: Vec<u32>,
     /// Flag to indicate if the JSON string of the notebook has a trailing newline.
@@ -128,7 +129,7 @@ impl Notebook {
 
         let mut contents = Vec::with_capacity(valid_code_cells.len());
         let mut current_offset = TextSize::from(0);
-        let mut cell_offsets = Vec::with_capacity(valid_code_cells.len());
+        let mut cell_offsets = CellOffsets::with_capacity(valid_code_cells.len());
         cell_offsets.push(TextSize::from(0));
 
         for &idx in &valid_code_cells {
@@ -324,9 +325,9 @@ impl Notebook {
         self.index.take().unwrap_or_else(|| self.build_index())
     }
 
-    /// Return the cell offsets for the concatenated source code corresponding
+    /// Return the [`CellOffsets`] for the concatenated source code corresponding
     /// the Jupyter notebook.
-    pub fn cell_offsets(&self) -> &[TextSize] {
+    pub fn cell_offsets(&self) -> &CellOffsets {
         &self.cell_offsets
     }
 
@@ -503,7 +504,7 @@ print("after empty cells")
             }
         );
         assert_eq!(
-            notebook.cell_offsets(),
+            notebook.cell_offsets().as_ref(),
             &[
                 0.into(),
                 90.into(),