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.
This commit is contained in:
Dhruv Manilawala 2023-11-22 09:27:00 -06:00 committed by GitHub
parent 0cb438dd65
commit 727e389cac
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 50 additions and 16 deletions

View file

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

View file

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

View file

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

View file

@ -1,4 +1,7 @@
use std::fmt; use std::fmt;
use std::ops::{Deref, DerefMut};
use ruff_text_size::TextSize;
use crate::schema::{Cell, SourceValue}; use crate::schema::{Cell, SourceValue};
@ -168,3 +171,34 @@ impl Cell {
lines.any(|line| line.trim_start().starts_with("%%")) 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
}
}

View file

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

View file

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