Support isort: skip, isort: on, and isort: off (#678)

Cargo.lock

@@ -1622,6 +1622,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "nohash-hasher"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451"
+
 [[package]]
 name = "nom"
 version = "5.1.2"
@@ -2239,6 +2245,7 @@ dependencies = [
  "anyhow",
  "assert_cmd",
  "bincode",
+ "bitflags",
  "cacache",
  "chrono",
  "clap 4.0.22",
@@ -2255,6 +2262,7 @@ dependencies = [
  "itertools",
  "libcst",
  "log",
+ "nohash-hasher",
  "notify",
  "num-bigint",
  "once_cell",

Cargo.toml

@@ -15,6 +15,7 @@ name = "ruff"
 [dependencies]
 anyhow = { version = "1.0.66" }
 bincode = { version = "1.3.3" }
+bitflags = { version = "1.3.2" }
 chrono = { version = "0.4.21" }
 clap = { version = "4.0.1", features = ["derive"] }
 colored = { version = "2.0.0" }
@@ -26,6 +27,7 @@ glob = { version = "0.3.0" }
 itertools = { version = "0.10.5" }
 libcst = { git = "https://github.com/charliermarsh/LibCST", rev = "a13ec97dd4eb925bde4d426c6e422582793b260c" }
 log = { version = "0.4.17" }
+nohash-hasher = { version = "0.2.0" }
 notify = { version = "4.0.17" }
 num-bigint = { version = "0.4.3" }
 once_cell = { version = "1.16.0" }

resources/test/fixtures/isort/skip.py

@@ -0,0 +1,10 @@
+# isort: off
+import sys
+import os
+import collections
+# isort: on
+
+import sys
+import os  # isort: skip
+import collections
+import abc

src/check_ast.rs

@@ -23,7 +23,6 @@ use crate::ast::{helpers, operations, visitor};
 use crate::autofix::fixer;
 use crate::checks::{Check, CheckCode, CheckKind};
 use crate::docstrings::definition::{Definition, DefinitionKind, Documentable};
-use crate::isort::track::ImportTracker;
 use crate::python::builtins::{BUILTINS, MAGIC_GLOBALS};
 use crate::python::future::ALL_FEATURE_NAMES;
 use crate::python::typing;
@@ -78,7 +77,6 @@ pub struct Checker<'a> {
     deferred_functions: Vec<(&'a Stmt, Vec<usize>, Vec<usize>, VisibleScope)>,
     deferred_lambdas: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
     deferred_assignments: Vec<usize>,
-    import_tracker: ImportTracker<'a>,
     // Internal, derivative state.
     visible_scope: VisibleScope,
     in_f_string: Option<Range>,
@@ -117,7 +115,6 @@ impl<'a> Checker<'a> {
             deferred_functions: Default::default(),
             deferred_lambdas: Default::default(),
             deferred_assignments: Default::default(),
-            import_tracker: ImportTracker::new(),
             // Internal, derivative state.
             visible_scope: VisibleScope {
                 modifier: Modifier::Module,
@@ -185,9 +182,6 @@ where
     'b: 'a,
 {
     fn visit_stmt(&mut self, stmt: &'b Stmt) {
-        // Call-through to any composed visitors.
-        self.import_tracker.visit_stmt(stmt);
-
         self.push_parent(stmt);
 
         // Track whether we've seen docstrings, non-imports, etc.
@@ -1674,9 +1668,6 @@ where
     }
 
     fn visit_excepthandler(&mut self, excepthandler: &'b Excepthandler) {
-        // Call-through to any composed visitors.
-        self.import_tracker.visit_excepthandler(excepthandler);
-
         match &excepthandler.node {
             ExcepthandlerKind::ExceptHandler { type_, name, .. } => {
                 if self.settings.enabled.contains(&CheckCode::E722) && type_.is_none() {

src/check_imports.rs

@@ -1,5 +1,6 @@
 //! Lint rules based on import analysis.
 
+use nohash_hasher::IntSet;
 use rustpython_parser::ast::Suite;
 
 use crate::ast::visitor::Visitor;
@@ -30,10 +31,11 @@ fn check_import_blocks(
 pub fn check_imports(
     python_ast: &Suite,
     locator: &SourceCodeLocator,
+    exclusions: &IntSet<usize>,
     settings: &Settings,
     autofix: &fixer::Mode,
 ) -> Vec<Check> {
-    let mut tracker = ImportTracker::new();
+    let mut tracker = ImportTracker::new(exclusions);
     for stmt in python_ast {
         tracker.visit_stmt(stmt);
     }

src/check_lines.rs

@@ -1,7 +1,6 @@
 //! Lint rules based on checking raw physical lines.
 
-use std::collections::BTreeMap;
+use nohash_hasher::IntMap;
-
 use once_cell::sync::Lazy;
 use regex::Regex;
 use rustpython_parser::ast::Location;
@@ -36,7 +35,7 @@ fn should_enforce_line_length(line: &str, length: usize, limit: usize) -> bool {
 pub fn check_lines(
     checks: &mut Vec<Check>,
     contents: &str,
-    noqa_line_for: &[usize],
+    noqa_line_for: &IntMap<usize, usize>,
     settings: &Settings,
     autofix: &fixer::Mode,
 ) {
@@ -44,8 +43,7 @@ pub fn check_lines(
     let enforce_line_too_long = settings.enabled.contains(&CheckCode::E501);
     let enforce_noqa = settings.enabled.contains(&CheckCode::M001);
 
-    let mut noqa_directives: BTreeMap<usize, (Directive, Vec<&str>)> = BTreeMap::new();
+    let mut noqa_directives: IntMap<usize, (Directive, Vec<&str>)> = IntMap::default();
-
     let mut line_checks = vec![];
     let mut ignored = vec![];
 
@@ -55,7 +53,7 @@ pub fn check_lines(
         // If there are newlines at the end of the file, they won't be represented in
         // `noqa_line_for`, so fallback to the current line.
         let noqa_lineno = noqa_line_for
-            .get(lineno)
+            .get(&lineno)
             .map(|lineno| lineno - 1)
             .unwrap_or(lineno);
 
@@ -153,7 +151,7 @@ pub fn check_lines(
         if let Some(line) = lines.last() {
             let lineno = lines.len() - 1;
             let noqa_lineno = noqa_line_for
-                .get(lineno)
+                .get(&lineno)
                 .map(|lineno| lineno - 1)
                 .unwrap_or(lineno);
 
@@ -257,6 +255,8 @@ pub fn check_lines(
 
 #[cfg(test)]
 mod tests {
+    use nohash_hasher::IntMap;
+
     use super::check_lines;
     use crate::autofix::fixer;
     use crate::checks::{Check, CheckCode};
@@ -265,7 +265,7 @@ mod tests {
     #[test]
     fn e501_non_ascii_char() {
         let line = "'\u{4e9c}' * 2"; // 7 in UTF-32, 9 in UTF-8.
-        let noqa_line_for: Vec<usize> = vec![1];
+        let noqa_line_for: IntMap<usize, usize> = Default::default();
         let check_with_max_line_length = |line_length: usize| {
             let mut checks: Vec<Check> = vec![];
             check_lines(

src/directives.rs

@@ -0,0 +1,206 @@
+//! Extract `# noqa` and `# isort: skip` directives from tokenized source.
+
+use bitflags::bitflags;
+use nohash_hasher::{IntMap, IntSet};
+use rustpython_ast::Location;
+use rustpython_parser::lexer::{LexResult, Tok};
+
+use crate::ast::types::Range;
+use crate::checks::LintSource;
+use crate::{Settings, SourceCodeLocator};
+
+bitflags! {
+    pub struct Flags: u32 {
+        const NOQA = 0b00000001;
+        const ISORT = 0b00000010;
+    }
+}
+
+impl Flags {
+    pub fn from_settings(settings: &Settings) -> Self {
+        if settings
+            .enabled
+            .iter()
+            .any(|check_code| matches!(check_code.lint_source(), LintSource::Imports))
+        {
+            Flags::NOQA | Flags::ISORT
+        } else {
+            Flags::NOQA
+        }
+    }
+}
+
+pub struct Directives {
+    pub noqa_line_for: IntMap<usize, usize>,
+    pub isort_exclusions: IntSet<usize>,
+}
+
+pub fn extract_directives(
+    lxr: &[LexResult],
+    locator: &SourceCodeLocator,
+    flags: &Flags,
+) -> Directives {
+    Directives {
+        noqa_line_for: if flags.contains(Flags::NOQA) {
+            extract_noqa_line_for(lxr)
+        } else {
+            Default::default()
+        },
+        isort_exclusions: if flags.contains(Flags::ISORT) {
+            extract_isort_exclusions(lxr, locator)
+        } else {
+            Default::default()
+        },
+    }
+}
+
+/// 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();
+    for (start, tok, end) in lxr.iter().flatten() {
+        if matches!(tok, Tok::EndOfFile) {
+            break;
+        }
+        // For multi-line strings, we expect `noqa` directives on the last line of the
+        // string.
+        if matches!(tok, Tok::String { .. }) && end.row() > start.row() {
+            for i in start.row()..end.row() {
+                noqa_line_for.insert(i, end.row());
+            }
+        }
+    }
+    noqa_line_for
+}
+
+/// Extract a set of lines over which to disable isort.
+pub fn extract_isort_exclusions(lxr: &[LexResult], locator: &SourceCodeLocator) -> IntSet<usize> {
+    let mut exclusions: IntSet<usize> = IntSet::default();
+    let mut off: Option<&Location> = None;
+    for (start, tok, end) in lxr.iter().flatten() {
+        // TODO(charlie): Modify RustPython to include the comment text in the token.
+        if matches!(tok, Tok::Comment) {
+            let comment_text = locator.slice_source_code_range(&Range {
+                location: *start,
+                end_location: *end,
+            });
+            if off.is_some() {
+                if comment_text == "# isort: on" {
+                    if let Some(start) = off {
+                        for row in start.row() + 1..=end.row() {
+                            exclusions.insert(row);
+                        }
+                    }
+                    off = None;
+                }
+            } else {
+                if comment_text.contains("isort: skip") || comment_text.contains("isort:skip") {
+                    exclusions.insert(start.row());
+                } else if comment_text == "# isort: off" {
+                    off = Some(start);
+                }
+            }
+        } else if matches!(tok, Tok::EndOfFile) {
+            if let Some(start) = off {
+                for row in start.row() + 1..=end.row() {
+                    exclusions.insert(row);
+                }
+            }
+            break;
+        }
+    }
+    exclusions
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use nohash_hasher::IntMap;
+    use rustpython_parser::lexer;
+    use rustpython_parser::lexer::LexResult;
+
+    use crate::directives::extract_noqa_line_for;
+
+    #[test]
+    fn extraction() -> Result<()> {
+        let empty: IntMap<usize, usize> = Default::default();
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = 1
+y = 2
+z = x + 1",
+        )
+        .collect();
+        assert_eq!(extract_noqa_line_for(&lxr), empty);
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "
+x = 1
+y = 2
+z = x + 1",
+        )
+        .collect();
+        assert_eq!(extract_noqa_line_for(&lxr), empty);
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = 1
+y = 2
+z = x + 1
+        ",
+        )
+        .collect();
+        assert_eq!(extract_noqa_line_for(&lxr), empty);
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = 1
+
+y = 2
+z = x + 1
+        ",
+        )
+        .collect();
+        assert_eq!(extract_noqa_line_for(&lxr), empty);
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = '''abc
+def
+ghi
+'''
+y = 2
+z = x + 1",
+        )
+        .collect();
+        assert_eq!(
+            extract_noqa_line_for(&lxr),
+            IntMap::from_iter([(1, 4), (2, 4), (3, 4)])
+        );
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = 1
+        y = '''abc
+        def
+        ghi
+        '''
+        z = 2",
+        )
+        .collect();
+        assert_eq!(
+            extract_noqa_line_for(&lxr),
+            IntMap::from_iter([(2, 5), (3, 5), (4, 5)])
+        );
+
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(
+            "x = 1
+        y = '''abc
+        def
+        ghi
+        '''",
+        )
+        .collect();
+        assert_eq!(
+            extract_noqa_line_for(&lxr),
+            IntMap::from_iter([(2, 5), (3, 5), (4, 5)])
+        );
+
+        Ok(())
+    }
+}

src/isort/mod.rs

@@ -208,17 +208,18 @@ mod tests {
     use crate::linter::test_path;
     use crate::Settings;
 
-    #[test_case(Path::new("reorder_within_section.py"))]
-    #[test_case(Path::new("no_reorder_within_section.py"))]
-    #[test_case(Path::new("separate_future_imports.py"))]
-    #[test_case(Path::new("separate_third_party_imports.py"))]
-    #[test_case(Path::new("separate_first_party_imports.py"))]
-    #[test_case(Path::new("deduplicate_imports.py"))]
     #[test_case(Path::new("combine_import_froms.py"))]
-    #[test_case(Path::new("preserve_indentation.py"))]
+    #[test_case(Path::new("deduplicate_imports.py"))]
     #[test_case(Path::new("fit_line_length.py"))]
     #[test_case(Path::new("import_from_after_import.py"))]
     #[test_case(Path::new("leading_prefix.py"))]
+    #[test_case(Path::new("no_reorder_within_section.py"))]
+    #[test_case(Path::new("preserve_indentation.py"))]
+    #[test_case(Path::new("reorder_within_section.py"))]
+    #[test_case(Path::new("separate_first_party_imports.py"))]
+    #[test_case(Path::new("separate_future_imports.py"))]
+    #[test_case(Path::new("separate_third_party_imports.py"))]
+    #[test_case(Path::new("skip.py"))]
     #[test_case(Path::new("trailing_suffix.py"))]
     fn isort(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.to_string_lossy());

src/isort/snapshots/ruff__isort__tests__skip.py.snap

@@ -0,0 +1,22 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 9
+    column: 0
+  end_location:
+    row: 11
+    column: 0
+  fix:
+    patch:
+      content: "import abc\nimport collections\n"
+      location:
+        row: 9
+        column: 0
+      end_location:
+        row: 11
+        column: 0
+    applied: false
+

src/isort/track.rs

@@ -1,3 +1,4 @@
+use nohash_hasher::IntSet;
 use rustpython_ast::{
     Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
     ExcepthandlerKind, Expr, ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, StmtKind,
@@ -8,16 +9,19 @@ use crate::ast::visitor::Visitor;
 
 #[derive(Debug)]
 pub struct ImportTracker<'a> {
-    pub blocks: Vec<Vec<&'a Stmt>>,
+    exclusions: &'a IntSet<usize>,
+    blocks: Vec<Vec<&'a Stmt>>,
 }
+
 impl<'a> ImportTracker<'a> {
-    pub fn new() -> Self {
+    pub fn new(exclusions: &'a IntSet<usize>) -> Self {
         Self {
+            exclusions,
             blocks: vec![vec![]],
         }
     }
 
-    fn add_import(&mut self, stmt: &'a Stmt) {
+    fn track_import(&mut self, stmt: &'a Stmt) {
         let index = self.blocks.len() - 1;
         self.blocks[index].push(stmt);
     }
@@ -43,8 +47,9 @@ where
         if matches!(
             stmt.node,
             StmtKind::Import { .. } | StmtKind::ImportFrom { .. }
-        ) {
+        ) && !self.exclusions.contains(&stmt.location.row())
-            self.add_import(stmt);
+        {
+            self.track_import(stmt);
         } else {
             self.finalize();
         }

src/lib.rs

@@ -25,6 +25,7 @@ pub mod checks_gen;
 pub mod cli;
 pub mod code_gen;
 mod cst;
+mod directives;
 mod docstrings;
 pub mod flake8_annotations;
 mod flake8_bugbear;
@@ -74,8 +75,12 @@ pub fn check(path: &Path, contents: &str, autofix: bool) -> Result<Vec<Check>> {
     // Initialize the SourceCodeLocator (which computes offsets lazily).
     let locator = SourceCodeLocator::new(contents);
 
-    // Determine the noqa line for every line in the source.
+    // Extract the `# noqa` and `# isort: skip` directives from the source.
-    let noqa_line_for = noqa::extract_noqa_line_for(&tokens);
+    let directives = directives::extract_directives(
+        &tokens,
+        &locator,
+        &directives::Flags::from_settings(&settings),
+    );
 
     // Generate checks.
     let checks = check_path(
@@ -83,7 +88,7 @@ pub fn check(path: &Path, contents: &str, autofix: bool) -> Result<Vec<Check>> {
         contents,
         tokens,
         &locator,
-        &noqa_line_for,
+        &directives,
         &settings,
         &if autofix { Mode::Generate } else { Mode::None },
     )?;

src/linter.rs

@@ -21,11 +21,12 @@ use crate::check_lines::check_lines;
 use crate::check_tokens::check_tokens;
 use crate::checks::{Check, CheckCode, CheckKind, LintSource};
 use crate::code_gen::SourceGenerator;
+use crate::directives::Directives;
 use crate::message::Message;
 use crate::noqa::add_noqa;
 use crate::settings::Settings;
 use crate::source_code_locator::SourceCodeLocator;
-use crate::{cache, fs, noqa};
+use crate::{cache, directives, fs};
 
 /// Collect tokens up to and including the first error.
 pub(crate) fn tokenize(contents: &str) -> Vec<LexResult> {
@@ -56,7 +57,7 @@ pub(crate) fn check_path(
     contents: &str,
     tokens: Vec<LexResult>,
     locator: &SourceCodeLocator,
-    noqa_line_for: &[usize],
+    directives: &Directives,
     settings: &Settings,
     autofix: &fixer::Mode,
 ) -> Result<Vec<Check>> {
@@ -88,7 +89,13 @@ pub(crate) fn check_path(
                     checks.extend(check_ast(&python_ast, locator, settings, autofix, path));
                 }
                 if use_imports {
-                    checks.extend(check_imports(&python_ast, locator, settings, autofix));
+                    checks.extend(check_imports(
+                        &python_ast,
+                        locator,
+                        &directives.isort_exclusions,
+                        settings,
+                        autofix,
+                    ));
                 }
             }
             Err(parse_error) => {
@@ -106,7 +113,13 @@ pub(crate) fn check_path(
     }
 
     // Run the lines-based checks.
-    check_lines(&mut checks, contents, noqa_line_for, settings, autofix);
+    check_lines(
+        &mut checks,
+        contents,
+        &directives.noqa_line_for,
+        settings,
+        autofix,
+    );
 
     // Create path ignores.
     if !checks.is_empty() && !settings.per_file_ignores.is_empty() {
@@ -134,8 +147,12 @@ pub fn lint_stdin(
     // Initialize the SourceCodeLocator (which computes offsets lazily).
     let locator = SourceCodeLocator::new(stdin);
 
-    // Determine the noqa line for every line in the source.
+    // Extract the `# noqa` and `# isort: skip` directives from the source.
-    let noqa_line_for = noqa::extract_noqa_line_for(&tokens);
+    let directives = directives::extract_directives(
+        &tokens,
+        &locator,
+        &directives::Flags::from_settings(settings),
+    );
 
     // Generate checks.
     let mut checks = check_path(
@@ -143,7 +160,7 @@ pub fn lint_stdin(
         stdin,
         tokens,
         &locator,
-        &noqa_line_for,
+        &directives,
         settings,
         autofix,
     )?;
@@ -188,8 +205,12 @@ pub fn lint_path(
     // Initialize the SourceCodeLocator (which computes offsets lazily).
     let locator = SourceCodeLocator::new(&contents);
 
-    // Determine the noqa line for every line in the source.
+    // Determine the noqa and isort exclusions.
-    let noqa_line_for = noqa::extract_noqa_line_for(&tokens);
+    let directives = directives::extract_directives(
+        &tokens,
+        &locator,
+        &directives::Flags::from_settings(settings),
+    );
 
     // Generate checks.
     let mut checks = check_path(
@@ -197,7 +218,7 @@ pub fn lint_path(
         &contents,
         tokens,
         &locator,
-        &noqa_line_for,
+        &directives,
         settings,
         autofix,
     )?;
@@ -230,8 +251,12 @@ pub fn add_noqa_to_path(path: &Path, settings: &Settings) -> Result<usize> {
     // Initialize the SourceCodeLocator (which computes offsets lazily).
     let locator = SourceCodeLocator::new(&contents);
 
-    // Determine the noqa line for every line in the source.
+    // Extract the `# noqa` and `# isort: skip` directives from the source.
-    let noqa_line_for = noqa::extract_noqa_line_for(&tokens);
+    let directives = directives::extract_directives(
+        &tokens,
+        &locator,
+        &directives::Flags::from_settings(settings),
+    );
 
     // Generate checks.
     let checks = check_path(
@@ -239,12 +264,12 @@ pub fn add_noqa_to_path(path: &Path, settings: &Settings) -> Result<usize> {
         &contents,
         tokens,
         &locator,
-        &noqa_line_for,
+        &directives,
         settings,
         &fixer::Mode::None,
     )?;
 
-    add_noqa(&checks, &contents, &noqa_line_for, path)
+    add_noqa(&checks, &contents, &directives.noqa_line_for, path)
 }
 
 pub fn autoformat_path(path: &Path) -> Result<()> {
@@ -268,13 +293,17 @@ pub fn test_path(path: &Path, settings: &Settings, autofix: &fixer::Mode) -> Res
     let contents = fs::read_file(path)?;
     let tokens: Vec<LexResult> = tokenize(&contents);
     let locator = SourceCodeLocator::new(&contents);
-    let noqa_line_for = noqa::extract_noqa_line_for(&tokens);
+    let directives = directives::extract_directives(
+        &tokens,
+        &locator,
+        &directives::Flags::from_settings(settings),
+    );
     check_path(
         path,
         &contents,
         tokens,
         &locator,
-        &noqa_line_for,
+        &directives,
         settings,
         autofix,
     )

src/noqa.rs

@@ -3,14 +3,14 @@ use std::fs;
 use std::path::Path;
 
 use anyhow::Result;
+use nohash_hasher::IntMap;
 use once_cell::sync::Lazy;
 use regex::Regex;
-use rustpython_parser::lexer::{LexResult, Tok};
 
 use crate::checks::{Check, CheckCode};
 
 static NO_QA_REGEX: Lazy<Regex> = Lazy::new(|| {
-    Regex::new(r"(?i)(?P<noqa>\s*# noqa(?::\s?(?P<codes>([A-Z]+[0-9]+(?:[,\s]+)?)+))?)")
+    Regex::new(r"(?P<noqa>\s*# noqa(?::\s?(?P<codes>([A-Z]+[0-9]+(?:[,\s]+)?)+))?)")
         .expect("Invalid regex")
 });
 static SPLIT_COMMA_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"[,\s]").expect("Invalid regex"));
@@ -43,30 +43,21 @@ pub fn extract_noqa_directive(line: &str) -> Directive {
     }
 }
 
-pub fn extract_noqa_line_for(lxr: &[LexResult]) -> Vec<usize> {
+pub fn add_noqa(
-    let mut noqa_line_for: Vec<usize> = vec![];
+    checks: &[Check],
-    for (start, tok, end) in lxr.iter().flatten() {
+    contents: &str,
-        if matches!(tok, Tok::EndOfFile) {
+    noqa_line_for: &IntMap<usize, usize>,
-            break;
+    path: &Path,
-        }
+) -> Result<usize> {
-        // For multi-line strings, we expect `noqa` directives on the last line of the
+    let (count, output) = add_noqa_inner(checks, contents, noqa_line_for)?;
-        // string. By definition, we can't have multiple multi-line strings on
+    fs::write(path, output)?;
-        // the same line, so we don't need to verify that we haven't already
+    Ok(count)
-        // traversed past the current line.
-        if matches!(tok, Tok::String { .. }) && end.row() > start.row() {
-            for i in (noqa_line_for.len())..(start.row() - 1) {
-                noqa_line_for.push(i + 1);
-            }
-            noqa_line_for.extend(vec![end.row(); (end.row() + 1) - start.row()]);
-        }
-    }
-    noqa_line_for
 }
 
 fn add_noqa_inner(
     checks: &[Check],
     contents: &str,
-    noqa_line_for: &[usize],
+    noqa_line_for: &IntMap<usize, usize>,
 ) -> Result<(usize, String)> {
     let lines: Vec<&str> = contents.lines().collect();
     let mut matches_by_line: BTreeMap<usize, BTreeSet<&CheckCode>> = BTreeMap::new();
@@ -82,7 +73,7 @@ fn add_noqa_inner(
         // If there are newlines at the end of the file, they won't be represented in
         // `noqa_line_for`, so fallback to the current line.
         let noqa_lineno = noqa_line_for
-            .get(lineno)
+            .get(&lineno)
             .map(|lineno| lineno - 1)
             .unwrap_or(lineno);
 
@@ -120,108 +111,20 @@ fn add_noqa_inner(
     Ok((count, output))
 }
 
-pub fn add_noqa(
-    checks: &[Check],
-    contents: &str,
-    noqa_line_for: &[usize],
-    path: &Path,
-) -> Result<usize> {
-    let (count, output) = add_noqa_inner(checks, contents, noqa_line_for)?;
-    fs::write(path, output)?;
-    Ok(count)
-}
-
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
     use rustpython_parser::ast::Location;
-    use rustpython_parser::lexer;
-    use rustpython_parser::lexer::LexResult;
 
     use crate::ast::types::Range;
     use crate::checks::{Check, CheckKind};
-    use crate::noqa::{add_noqa_inner, extract_noqa_line_for};
+    use crate::noqa::add_noqa_inner;
-
-    #[test]
-    fn extraction() -> Result<()> {
-        let empty: Vec<usize> = Default::default();
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = 1
-y = 2
-z = x + 1",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), empty);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "
-x = 1
-y = 2
-z = x + 1",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), empty);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = 1
-y = 2
-z = x + 1
-        ",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), empty);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = 1
-
-y = 2
-z = x + 1
-        ",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), empty);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = '''abc
-def
-ghi
-'''
-y = 2
-z = x + 1",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), vec![4, 4, 4, 4]);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = 1
-y = '''abc
-def
-ghi
-'''
-z = 2",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), vec![1, 5, 5, 5, 5]);
-
-        let lxr: Vec<LexResult> = lexer::make_tokenizer(
-            "x = 1
-y = '''abc
-def
-ghi
-'''",
-        )
-        .collect();
-        assert_eq!(extract_noqa_line_for(&lxr), vec![1, 5, 5, 5, 5]);
-
-        Ok(())
-    }
 
     #[test]
     fn modification() -> Result<()> {
         let checks = vec![];
         let contents = "x = 1";
-        let noqa_line_for = vec![1];
+        let noqa_line_for = Default::default();
         let (count, output) = add_noqa_inner(&checks, contents, &noqa_line_for)?;
         assert_eq!(count, 0);
         assert_eq!(output.trim(), contents.trim());
@@ -234,7 +137,7 @@ ghi
             },
         )];
         let contents = "x = 1";
-        let noqa_line_for = vec![1];
+        let noqa_line_for = Default::default();
         let (count, output) = add_noqa_inner(&checks, contents, &noqa_line_for)?;
         assert_eq!(count, 1);
         assert_eq!(output.trim(), "x = 1  # noqa: F841".trim());
@@ -256,7 +159,7 @@ ghi
             ),
         ];
         let contents = "x = 1  # noqa: E741";
-        let noqa_line_for = vec![1];
+        let noqa_line_for = Default::default();
         let (count, output) = add_noqa_inner(&checks, contents, &noqa_line_for)?;
         assert_eq!(count, 1);
         assert_eq!(output.trim(), "x = 1  # noqa: E741, F841".trim());
@@ -278,7 +181,7 @@ ghi
             ),
         ];
         let contents = "x = 1  # noqa";
-        let noqa_line_for = vec![1];
+        let noqa_line_for = Default::default();
         let (count, output) = add_noqa_inner(&checks, contents, &noqa_line_for)?;
         assert_eq!(count, 1);
         assert_eq!(output.trim(), "x = 1  # noqa: E741, F841".trim());

src/source_code_locator.rs

@@ -17,7 +17,7 @@ impl<'a> SourceCodeLocator<'a> {
     pub fn new(contents: &'a str) -> Self {
         SourceCodeLocator {
             contents,
-            rope: OnceCell::new(),
+            rope: Default::default(),
         }
     }