Enable noqa directives on logical lines

Cargo.toml

@@ -26,7 +26,7 @@ once_cell = { version = "1.13.1" }
 path-absolutize = { version = "3.0.13", features = ["once_cell_cache"] }
 rayon = { version = "1.5.3" }
 regex = { version = "1.6.0" }
-rustpython-parser = { features = ["lalrpop"], git = "https://github.com/charliermarsh/RustPython.git", rev = "7d21c6923a506e79cc041708d83cef925efd33f4" }
+rustpython-parser = { features = ["lalrpop"], git = "https://github.com/charliermarsh/RustPython.git", rev = "966a80597d626a9a47eaec78471164422d341453" }
 serde = { version = "1.0.143", features = ["derive"] }
 serde_json = { version = "1.0.83" }
 toml = { version = "0.5.9" }

resources/test/fixtures/E501.py

@@ -5,5 +5,12 @@
 Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
 """
 
+_ = """Lorem ipsum dolor sit amet.
+
+    https://github.com/PyCQA/pycodestyle/pull/258/files#diff-841c622497a8033d10152bfdfb15b20b92437ecdea21a260944ea86b77b51533
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+"""  # noqa: E501
+
 _ = "---------------------------------------------------------------------------AAAAAAA"
 _ = "---------------------------------------------------------------------------亜亜亜亜亜亜亜"

resources/test/fixtures/M001.py

@@ -14,3 +14,25 @@ def f() -> None:
 
     # Invalid
     d = 1  # noqa: F841, E501
+
+
+_ = """Lorem ipsum dolor sit amet.
+
+    https://github.com/PyCQA/pycodestyle/pull/258/files#diff-841c622497a8033d10152bfdfb15b20b92437ecdea21a260944ea86b77b51533
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+"""  # noqa: E501
+
+_ = """Lorem ipsum dolor sit amet.
+
+    https://github.com/PyCQA/pycodestyle/pull/258/files#diff-841c622497a8033d10152bfdfb15b20b92437ecdea21a260944ea86b77b51533
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+"""  # noqa: E501, F841
+
+_ = """Lorem ipsum dolor sit amet.
+
+    https://github.com/PyCQA/pycodestyle/pull/258/files#diff-841c622497a8033d10152bfdfb15b20b92437ecdea21a260944ea86b77b51533
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.  # noqa: E501
+"""  # noqa: E501

src/check_lines.rs

@@ -1,4 +1,5 @@
-use once_cell::sync::Lazy;
+use std::collections::BTreeMap;
+
 use rustpython_parser::ast::Location;
 
 use crate::checks::{Check, CheckCode, CheckKind};
@@ -22,32 +23,84 @@ fn should_enforce_line_length(line: &str, length: usize, limit: usize) -> bool {
     }
 }
 
-pub fn check_lines(checks: &mut Vec<Check>, contents: &str, settings: &Settings) {
+pub fn check_lines(
+    checks: &mut Vec<Check>,
+    contents: &str,
+    noqa_line_for: &[usize],
+    settings: &Settings,
+) {
     let enforce_line_too_long = settings.select.contains(&CheckCode::E501);
     let enforce_noqa = settings.select.contains(&CheckCode::M001);
 
+    let mut noqa_directives: BTreeMap<usize, (Directive, Vec<&str>)> = BTreeMap::new();
+
     let mut line_checks = vec![];
     let mut ignored = vec![];
-    for (row, line) in contents.lines().enumerate() {
-        let noqa_directive = Lazy::new(|| noqa::extract_noqa_directive(line));
-        let mut line_ignored: Vec<&str> = vec![];
+
+    let lines: Vec<&str> = contents.lines().collect();
+    for (lineno, line) in lines.iter().enumerate() {
+        let mut did_insert: bool = false;
+
+        // Grab the noqa (logical) line number for the current (physical) line.
+        // 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)
+            .map(|lineno| lineno - 1)
+            .unwrap_or(lineno);
+
+        if enforce_noqa && !did_insert {
+            // Try the current physical line.
+            noqa_directives
+                .entry(lineno)
+                .or_insert_with(|| (noqa::extract_noqa_directive(lines[lineno]), vec![]));
+            // Try the current logical line.
+            if lineno != noqa_lineno {
+                noqa_directives
+                    .entry(noqa_lineno)
+                    .or_insert_with(|| (noqa::extract_noqa_directive(lines[noqa_lineno]), vec![]));
+            }
+            did_insert = true;
+        }
 
         // Remove any ignored checks.
         // TODO(charlie): Only validate checks for the current line.
         for (index, check) in checks.iter().enumerate() {
-            if check.location.row() == row + 1 {
-                match &*noqa_directive {
-                    Directive::All(_) => {
-                        line_ignored.push(check.kind.code().as_str());
+            if check.location.row() == lineno + 1 {
+                if !did_insert {
+                    // Try the current physical line.
+                    noqa_directives
+                        .entry(lineno)
+                        .or_insert_with(|| (noqa::extract_noqa_directive(lines[lineno]), vec![]));
+                    // Try the current logical line.
+                    if lineno != noqa_lineno {
+                        noqa_directives.entry(noqa_lineno).or_insert_with(|| {
+                            (noqa::extract_noqa_directive(lines[noqa_lineno]), vec![])
+                        });
+                    }
+                    did_insert = true;
+                }
+
+                let noqa = if lineno != noqa_lineno
+                    && matches!(noqa_directives.get(&lineno).unwrap(), (Directive::None, _))
+                {
+                    noqa_directives.get_mut(&noqa_lineno).unwrap()
+                } else {
+                    noqa_directives.get_mut(&lineno).unwrap()
+                };
+
+                match noqa {
+                    (Directive::All(_), matches) => {
+                        matches.push(check.kind.code().as_str());
                         ignored.push(index)
                     }
-                    Directive::Codes(_, codes) => {
+                    (Directive::Codes(_, codes), matches) => {
                         if codes.contains(&check.kind.code().as_str()) {
-                            line_ignored.push(check.kind.code().as_str());
+                            matches.push(check.kind.code().as_str());
                             ignored.push(index);
                         }
                     }
-                    Directive::None => {}
+                    (Directive::None, _) => {}
                 }
             }
         }
@@ -56,40 +109,64 @@ pub fn check_lines(checks: &mut Vec<Check>, contents: &str, settings: &Settings)
         if enforce_line_too_long {
             let line_length = line.chars().count();
             if should_enforce_line_length(line, line_length, settings.line_length) {
+                if !did_insert {
+                    // Try the current physical line.
+                    noqa_directives
+                        .entry(lineno)
+                        .or_insert_with(|| (noqa::extract_noqa_directive(lines[lineno]), vec![]));
+                    // Try the current logical line.
+                    if lineno != noqa_lineno {
+                        noqa_directives.entry(noqa_lineno).or_insert_with(|| {
+                            (noqa::extract_noqa_directive(lines[noqa_lineno]), vec![])
+                        });
+                    }
+                }
+
+                let noqa = if lineno != noqa_lineno
+                    && matches!(noqa_directives.get(&lineno).unwrap(), (Directive::None, _))
+                {
+                    noqa_directives.get_mut(&noqa_lineno).unwrap()
+                } else {
+                    noqa_directives.get_mut(&lineno).unwrap()
+                };
+
                 let check = Check::new(
                     CheckKind::LineTooLong(line_length, settings.line_length),
-                    Location::new(row + 1, settings.line_length + 1),
+                    Location::new(lineno + 1, settings.line_length + 1),
                 );
-                match &*noqa_directive {
-                    Directive::All(_) => {
-                        line_ignored.push(check.kind.code().as_str());
+
+                match noqa {
+                    (Directive::All(_), matches) => {
+                        matches.push(check.kind.code().as_str());
                     }
-                    Directive::Codes(_, codes) => {
+                    (Directive::Codes(_, codes), matches) => {
                         if codes.contains(&check.kind.code().as_str()) {
-                            line_ignored.push(check.kind.code().as_str());
+                            matches.push(check.kind.code().as_str());
                         } else {
                             line_checks.push(check);
                         }
                     }
-                    Directive::None => line_checks.push(check),
+                    (Directive::None, _) => line_checks.push(check),
                 }
             }
         }
+    }
 
-        // Enforce that the noqa was actually used.
-        if enforce_noqa {
-            match &*noqa_directive {
+    // Enforce that the noqa directive was actually used.
+    if enforce_noqa {
+        for (row, (directive, matches)) in noqa_directives {
+            match directive {
                 Directive::All(column) => {
-                    if line_ignored.is_empty() {
+                    if matches.is_empty() {
                         line_checks.push(Check::new(
                             CheckKind::UnusedNOQA(None),
                             Location::new(row + 1, column + 1),
                         ));
                     }
                 }
                 Directive::Codes(column, codes) => {
-                    for code in codes {
-                        if !line_ignored.contains(code) {
+                    for code in &codes {
+                        if !matches.contains(code) {
                             line_checks.push(Check::new(
                                 CheckKind::UnusedNOQA(Some(code.to_string())),
                                 Location::new(row + 1, column + 1),
@@ -101,6 +178,7 @@ pub fn check_lines(checks: &mut Vec<Check>, contents: &str, settings: &Settings)
             }
         }
     }
+
     ignored.sort();
     for index in ignored.iter().rev() {
         checks.swap_remove(*index);
@@ -118,6 +196,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 check_with_max_line_length = |line_length: usize| {
             let mut checks: Vec<Check> = vec![];
             let settings = Settings {
@@ -128,7 +207,8 @@ mod tests {
                 extend_exclude: vec![],
                 select: BTreeSet::from_iter(vec![CheckCode::E501]),
             };
-            check_lines(&mut checks, line, &settings);
+
+            check_lines(&mut checks, line, &noqa_line_for, &settings);
             return checks;
         };
         assert!(!check_with_max_line_length(6).is_empty());

src/linter.rs

@@ -2,7 +2,8 @@ use std::path::Path;
 
 use anyhow::Result;
 use log::debug;
-use rustpython_parser::parser;
+use rustpython_parser::lexer::LexResult;
+use rustpython_parser::{lexer, parser};
 
 use crate::autofix::fixer;
 use crate::autofix::fixer::fix_file;
@@ -11,7 +12,7 @@ use crate::check_lines::check_lines;
 use crate::checks::{Check, CheckCode, CheckKind, LintSource};
 use crate::message::Message;
 use crate::settings::Settings;
-use crate::{cache, fs};
+use crate::{cache, fs, noqa};
 
 fn check_path(
     path: &Path,
@@ -22,13 +23,19 @@ fn check_path(
     // Aggregate all checks.
     let mut checks: Vec<Check> = vec![];
 
+    // Tokenize once.
+    let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
+
+    // Determine the noqa line for every line in the source.
+    let noqa_line_for = noqa::extract_line_map(&lxr);
+
     // Run the AST-based checks.
     if settings
         .select
         .iter()
         .any(|check_code| matches!(check_code.lint_source(), LintSource::AST))
     {
-        match parser::parse_program(contents, "<filename>") {
+        match parser::parse_program_tokens(lxr, "<filename>") {
             Ok(python_ast) => {
                 checks.extend(check_ast(&python_ast, contents, settings, autofix, path))
             }
@@ -44,7 +51,7 @@ fn check_path(
     }
 
     // Run the lines-based checks.
-    check_lines(&mut checks, contents, settings);
+    check_lines(&mut checks, contents, &noqa_line_for, settings);
 
     checks
 }
@@ -582,7 +589,7 @@ mod tests {
     fn m001() -> Result<()> {
         let mut checks = check_path(
             Path::new("./resources/test/fixtures/M001.py"),
-            &settings::Settings::for_rules(vec![CheckCode::M001, CheckCode::F841]),
+            &settings::Settings::for_rules(vec![CheckCode::M001, CheckCode::E501, CheckCode::F841]),
             &fixer::Mode::Generate,
         )?;
         checks.sort_by_key(|check| check.location);