ruff_notebook/
index.rs

1use serde::{Deserialize, Serialize};
2
3use ruff_source_file::{LineColumn, OneIndexed, SourceLocation};
4
5/// Jupyter Notebook indexing table
6///
7/// When we lint a jupyter notebook, we have to translate the row/column based on
8/// [`ruff_text_size::TextSize`] to jupyter notebook cell/row/column.
9#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
10pub struct NotebookIndex {
11    /// Stores the starting row and the absolute cell index for every Python (valid) cell.
12    ///
13    /// The index in this vector corresponds to the Python cell index (valid cell index).
14    pub(super) cell_starts: Vec<CellStart>,
15}
16
17impl NotebookIndex {
18    fn find_cell(&self, row: OneIndexed) -> Option<CellStart> {
19        match self
20            .cell_starts
21            .binary_search_by_key(&row, |start| start.start_row)
22        {
23            Ok(cell_index) => Some(self.cell_starts[cell_index]),
24            Err(insertion_point) => Some(self.cell_starts[insertion_point.checked_sub(1)?]),
25        }
26    }
27
28    /// Returns the (raw) cell number (1-based) for the given row (1-based).
29    pub fn cell(&self, row: OneIndexed) -> Option<OneIndexed> {
30        self.find_cell(row).map(|start| start.raw_cell_index)
31    }
32
33    /// Returns the row number (1-based) in the cell (1-based) for the
34    /// given row (1-based).
35    pub fn cell_row(&self, row: OneIndexed) -> Option<OneIndexed> {
36        self.find_cell(row)
37            .map(|start| OneIndexed::from_zero_indexed(row.get() - start.start_row.get()))
38    }
39
40    /// Returns an iterator over the starting rows of each cell (1-based).
41    ///
42    /// This yields one entry per Python cell (skipping over Markdown cell).
43    pub fn iter(&self) -> impl Iterator<Item = CellStart> + '_ {
44        self.cell_starts.iter().copied()
45    }
46
47    /// Translates the given [`LineColumn`] based on the indexing table.
48    ///
49    /// This will translate the row/column in the concatenated source code
50    /// to the row/column in the Jupyter Notebook cell.
51    pub fn translate_line_column(&self, source_location: &LineColumn) -> LineColumn {
52        LineColumn {
53            line: self
54                .cell_row(source_location.line)
55                .unwrap_or(OneIndexed::MIN),
56            column: source_location.column,
57        }
58    }
59
60    /// Translates the given [`SourceLocation`] based on the indexing table.
61    ///
62    /// This will translate the line/character in the concatenated source code
63    /// to the line/character in the Jupyter Notebook cell.
64    pub fn translate_source_location(&self, source_location: &SourceLocation) -> SourceLocation {
65        SourceLocation {
66            line: self
67                .cell_row(source_location.line)
68                .unwrap_or(OneIndexed::MIN),
69            character_offset: source_location.character_offset,
70        }
71    }
72}
73
74#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)]
75pub struct CellStart {
76    /// The row in the concatenated notebook source code at which
77    /// this cell starts.
78    pub(super) start_row: OneIndexed,
79
80    /// The absolute index of this cell in the notebook.
81    pub(super) raw_cell_index: OneIndexed,
82}
83
84impl CellStart {
85    pub fn start_row(&self) -> OneIndexed {
86        self.start_row
87    }
88
89    pub fn cell_index(&self) -> OneIndexed {
90        self.raw_cell_index
91    }
92}
ruff_notebook/index.rs

ruff_notebook/
index.rs
