Add a dedicated token indexer for continuations and comments (#1886)

The primary motivation is that we can now robustly detect `\` continuations due to the addition of `Tok::NonLogicalNewline`. This PR generalizes the approach we took to comments (track all lines that contain any comments), and applies it to continuations too.

src/ast/helpers.rs

@@ -13,7 +13,7 @@ use rustpython_parser::token::StringKind;
 
 use crate::ast::types::{Binding, BindingKind, Range};
 use crate::checkers::ast::Checker;
-use crate::source_code::{Generator, Locator, Stylist};
+use crate::source_code::{Generator, Indexer, Locator, Stylist};
 
 /// Create an `Expr` with default location from an `ExprKind`.
 pub fn create_expr(node: ExprKind) -> Expr {
@@ -601,27 +601,6 @@ pub fn else_range(stmt: &Stmt, locator: &Locator) -> Option<Range> {
     }
 }
 
-/// Return `true` if a `Stmt` appears to be part of a multi-statement line, with
-/// other statements preceding it.
-pub fn preceded_by_continuation(stmt: &Stmt, locator: &Locator) -> bool {
-    // Does the previous line end in a continuation? This will have a specific
-    // false-positive, which is that if the previous line ends in a comment, it
-    // will be treated as a continuation. So we should only use this information to
-    // make conservative choices.
-    // TODO(charlie): Come up with a more robust strategy.
-    if stmt.location.row() > 1 {
-        let range = Range::new(
-            Location::new(stmt.location.row() - 1, 0),
-            Location::new(stmt.location.row(), 0),
-        );
-        let line = locator.slice_source_code_range(&range);
-        if line.trim_end().ends_with('\\') {
-            return true;
-        }
-    }
-    false
-}
-
 /// Return the `Range` of the first `Tok::Colon` token in a `Range`.
 pub fn first_colon_range(range: Range, locator: &Locator) -> Option<Range> {
     let contents = locator.slice_source_code_range(&range);
@@ -637,8 +616,17 @@ pub fn first_colon_range(range: Range, locator: &Locator) -> Option<Range> {
 
 /// Return `true` if a `Stmt` appears to be part of a multi-statement line, with
 /// other statements preceding it.
-pub fn preceded_by_multi_statement_line(stmt: &Stmt, locator: &Locator) -> bool {
-    match_leading_content(stmt, locator) || preceded_by_continuation(stmt, locator)
+pub fn preceded_by_continuation(stmt: &Stmt, indexer: &Indexer) -> bool {
+    stmt.location.row() > 1
+        && indexer
+            .continuation_lines()
+            .contains(&(stmt.location.row() - 1))
+}
+
+/// Return `true` if a `Stmt` appears to be part of a multi-statement line, with
+/// other statements preceding it.
+pub fn preceded_by_multi_statement_line(stmt: &Stmt, locator: &Locator, indexer: &Indexer) -> bool {
+    match_leading_content(stmt, locator) || preceded_by_continuation(stmt, indexer)
 }
 
 /// Return `true` if a `Stmt` appears to be part of a multi-statement line, with

src/autofix/helpers.rs

@@ -12,7 +12,7 @@ use crate::ast::whitespace::LinesWithTrailingNewline;
 use crate::cst::helpers::compose_module_path;
 use crate::cst::matchers::match_module;
 use crate::fix::Fix;
-use crate::source_code::Locator;
+use crate::source_code::{Indexer, Locator};
 
 /// Determine if a body contains only a single statement, taking into account
 /// deleted.
@@ -156,6 +156,7 @@ pub fn delete_stmt(
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
     locator: &Locator,
+    indexer: &Indexer,
 ) -> Result<Fix> {
     if parent
         .map(|parent| is_lone_child(stmt, parent, deleted))
@@ -175,7 +176,7 @@ pub fn delete_stmt(
             Fix::deletion(stmt.location, next)
         } else if helpers::match_leading_content(stmt, locator) {
             Fix::deletion(stmt.location, stmt.end_location.unwrap())
-        } else if helpers::preceded_by_continuation(stmt, locator) {
+        } else if helpers::preceded_by_continuation(stmt, indexer) {
             if is_end_of_file(stmt, locator) && stmt.location.column() == 0 {
                 // Special-case: a file can't end in a continuation.
                 Fix::replacement("\n".to_string(), stmt.location, stmt.end_location.unwrap())
@@ -198,6 +199,7 @@ pub fn remove_unused_imports<'a>(
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
     locator: &Locator,
+    indexer: &Indexer,
 ) -> Result<Fix> {
     let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
     let mut tree = match_module(&module_text)?;
@@ -235,7 +237,7 @@ pub fn remove_unused_imports<'a>(
                 if !found_star {
                     bail!("Expected \'*\' for unused import");
                 }
-                return delete_stmt(stmt, parent, deleted, locator);
+                return delete_stmt(stmt, parent, deleted, locator, indexer);
             } else {
                 bail!("Expected: ImportNames::Aliases | ImportNames::Star");
             }
@@ -296,7 +298,7 @@ pub fn remove_unused_imports<'a>(
     }
 
     if aliases.is_empty() {
-        delete_stmt(stmt, parent, deleted, locator)
+        delete_stmt(stmt, parent, deleted, locator, indexer)
     } else {
         let mut state = CodegenState::default();
         tree.codegen(&mut state);

src/checkers/ast.rs

@@ -39,7 +39,7 @@ use crate::rules::{
 };
 use crate::settings::types::PythonVersion;
 use crate::settings::{flags, Settings};
-use crate::source_code::{Locator, Stylist};
+use crate::source_code::{Indexer, Locator, Stylist};
 use crate::violations::DeferralKeyword;
 use crate::visibility::{module_visibility, transition_scope, Modifier, Visibility, VisibleScope};
 use crate::{autofix, docstrings, noqa, violations, visibility};
@@ -57,7 +57,8 @@ pub struct Checker<'a> {
     pub(crate) settings: &'a Settings,
     pub(crate) noqa_line_for: &'a IntMap<usize, usize>,
     pub(crate) locator: &'a Locator<'a>,
-    pub(crate) style: &'a Stylist<'a>,
+    pub(crate) stylist: &'a Stylist<'a>,
+    pub(crate) indexer: &'a Indexer,
     // Computed diagnostics.
     pub(crate) diagnostics: Vec<Diagnostic>,
     // Function and class definition tracking (e.g., for docstring enforcement).
@@ -98,6 +99,7 @@ pub struct Checker<'a> {
 }
 
 impl<'a> Checker<'a> {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         settings: &'a Settings,
         noqa_line_for: &'a IntMap<usize, usize>,
@@ -106,6 +108,7 @@ impl<'a> Checker<'a> {
         path: &'a Path,
         locator: &'a Locator,
         style: &'a Stylist,
+        indexer: &'a Indexer,
     ) -> Checker<'a> {
         Checker {
             settings,
@@ -114,7 +117,8 @@ impl<'a> Checker<'a> {
             noqa,
             path,
             locator,
-            style,
+            stylist: style,
+            indexer,
             diagnostics: vec![],
             definitions: vec![],
             deletions: FxHashSet::default(),
@@ -4001,6 +4005,7 @@ impl<'a> Checker<'a> {
                             parent,
                             &deleted,
                             self.locator,
+                            self.indexer,
                         ) {
                             Ok(fix) => {
                                 if fix.content.is_empty() || fix.content == "pass" {
@@ -4296,6 +4301,7 @@ pub fn check_ast(
     python_ast: &Suite,
     locator: &Locator,
     stylist: &Stylist,
+    indexer: &Indexer,
     noqa_line_for: &IntMap<usize, usize>,
     settings: &Settings,
     autofix: flags::Autofix,
@@ -4310,6 +4316,7 @@ pub fn check_ast(
         path,
         locator,
         stylist,
+        indexer,
     );
     checker.push_scope(Scope::new(ScopeKind::Module));
     checker.bind_builtins();

src/checkers/imports.rs

@@ -10,12 +10,13 @@ use crate::registry::{Diagnostic, RuleCode};
 use crate::rules::isort;
 use crate::rules::isort::track::{Block, ImportTracker};
 use crate::settings::{flags, Settings};
-use crate::source_code::{Locator, Stylist};
+use crate::source_code::{Indexer, Locator, Stylist};
 
 #[allow(clippy::too_many_arguments)]
 pub fn check_imports(
     python_ast: &Suite,
     locator: &Locator,
+    indexer: &Indexer,
     directives: &IsortDirectives,
     settings: &Settings,
     stylist: &Stylist,
@@ -39,7 +40,7 @@ pub fn check_imports(
         for block in &blocks {
             if !block.imports.is_empty() {
                 if let Some(diagnostic) = isort::rules::organize_imports(
-                    block, locator, settings, stylist, autofix, package,
+                    block, locator, indexer, settings, stylist, autofix, package,
                 ) {
                     diagnostics.push(diagnostic);
                 }

src/directives.rs

@@ -37,14 +37,12 @@ pub struct IsortDirectives {
 }
 
 pub struct Directives {
-    pub commented_lines: Vec<usize>,
     pub noqa_line_for: IntMap<usize, usize>,
     pub isort: IsortDirectives,
 }
 
 pub fn extract_directives(lxr: &[LexResult], flags: Flags) -> Directives {
     Directives {
-        commented_lines: extract_commented_lines(lxr),
         noqa_line_for: if flags.contains(Flags::NOQA) {
             extract_noqa_line_for(lxr)
         } else {
@@ -58,16 +56,6 @@ pub fn extract_directives(lxr: &[LexResult], flags: Flags) -> Directives {
     }
 }
 
-pub fn extract_commented_lines(lxr: &[LexResult]) -> Vec<usize> {
-    let mut commented_lines = Vec::new();
-    for (start, tok, ..) in lxr.iter().flatten() {
-        if matches!(tok, Tok::Comment(_)) {
-            commented_lines.push(start.row());
-        }
-    }
-    commented_lines
-}
-
 /// Extract a mapping from logical line to noqa line.
 pub fn extract_noqa_line_for(lxr: &[LexResult]) -> IntMap<usize, usize> {
     let mut noqa_line_for: IntMap<usize, usize> = IntMap::default();

src/lib_native.rs

@@ -10,7 +10,7 @@ use crate::resolver::Relativity;
 use crate::rustpython_helpers::tokenize;
 use crate::settings::configuration::Configuration;
 use crate::settings::{flags, pyproject, Settings};
-use crate::source_code::{Locator, Stylist};
+use crate::source_code::{Indexer, Locator, Stylist};
 use crate::{directives, packaging, resolver};
 
 /// Load the relevant `Settings` for a given `Path`.
@@ -44,6 +44,9 @@ pub fn check(path: &Path, contents: &str, autofix: bool) -> Result<Vec<Diagnosti
     // Detect the current code style (lazily).
     let stylist = Stylist::from_contents(contents, &locator);
 
+    // Extra indices from the code.
+    let indexer: Indexer = tokens.as_slice().into();
+
     // Extract the `# noqa` and `# isort: skip` directives from the source.
     let directives =
         directives::extract_directives(&tokens, directives::Flags::from_settings(&settings));
@@ -56,6 +59,7 @@ pub fn check(path: &Path, contents: &str, autofix: bool) -> Result<Vec<Diagnosti
         tokens,
         &locator,
         &stylist,
+        &indexer,
         &directives,
         &settings,
         autofix.into(),

src/lib_wasm.rs

@@ -18,7 +18,7 @@ use crate::settings::configuration::Configuration;
 use crate::settings::options::Options;
 use crate::settings::types::PythonVersion;
 use crate::settings::{flags, Settings};
-use crate::source_code::{Locator, Stylist};
+use crate::source_code::{Indexer, Locator, Stylist};
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
@@ -157,6 +157,9 @@ pub fn check(contents: &str, options: JsValue) -> Result<JsValue, JsValue> {
     // Detect the current code style (lazily).
     let stylist = Stylist::from_contents(contents, &locator);
 
+    // Extra indices from the code.
+    let indexer: Indexer = tokens.as_slice().into();
+
     // Extract the `# noqa` and `# isort: skip` directives from the source.
     let directives = directives::extract_directives(&tokens, directives::Flags::empty());
 
@@ -168,6 +171,7 @@ pub fn check(contents: &str, options: JsValue) -> Result<JsValue, JsValue> {
         tokens,
         &locator,
         &stylist,
+        &indexer,
         &directives,
         &settings,
         flags::Autofix::Enabled,