From 727e389cac222e64c6f89ea1341f99ee54109b9d Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 22 Nov 2023 09:27:00 -0600 Subject: [PATCH] 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 | 6 ++-- crates/ruff_linter/src/linter.rs | 4 ++- crates/ruff_linter/src/rules/isort/block.rs | 10 ++---- crates/ruff_notebook/src/cell.rs | 34 +++++++++++++++++++++ crates/ruff_notebook/src/lib.rs | 1 + crates/ruff_notebook/src/notebook.rs | 11 ++++--- 6 files changed, 50 insertions(+), 16 deletions(-) diff --git a/crates/ruff_linter/src/checkers/imports.rs b/crates/ruff_linter/src/checkers/imports.rs index 22cbe662f5..5ecb477d85 100644 --- a/crates/ruff_linter/src/checkers/imports.rs +++ b/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 { 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, Option) { // 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 }; diff --git a/crates/ruff_linter/src/linter.rs b/crates/ruff_linter/src/linter.rs index 328c6142ae..8ebac983ec 100644 --- a/crates/ruff_linter/src/linter.rs +++ b/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); diff --git a/crates/ruff_linter/src/rules/isort/block.rs b/crates/ruff_linter/src/rules/isort/block.rs index 4f24a1c1d2..8389fa86a2 100644 --- a/crates/ruff_linter/src/rules/isort/block.rs +++ b/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()), } } diff --git a/crates/ruff_notebook/src/cell.rs b/crates/ruff_notebook/src/cell.rs index 2251df1331..c992febd89 100644 --- a/crates/ruff_notebook/src/cell.rs +++ b/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); + +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 + } +} diff --git a/crates/ruff_notebook/src/lib.rs b/crates/ruff_notebook/src/lib.rs index 03271682f8..e95c483a2f 100644 --- a/crates/ruff_notebook/src/lib.rs +++ b/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::*; diff --git a/crates/ruff_notebook/src/notebook.rs b/crates/ruff_notebook/src/notebook.rs index 3deb08be83..a6714fa12b 100644 --- a/crates/ruff_notebook/src/notebook.rs +++ b/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, + cell_offsets: CellOffsets, /// The cell index of all valid code cells in the notebook. valid_code_cells: Vec, /// 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(),