From 2ea2495111279cd624a9f35007c5ef158bf1e4f9 Mon Sep 17 00:00:00 2001
From: augustelalande <auguste.lalande@gmail.com>
Date: Wed, 24 Apr 2024 14:01:45 -0400
Subject: [PATCH] improve handling of EOF in noqa

---
 crates/ruff_linter/src/noqa.rs                | 46 ++++++++++++-------
 .../rules/pygrep_hooks/rules/blanket_noqa.rs  | 10 +---
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/crates/ruff_linter/src/noqa.rs b/crates/ruff_linter/src/noqa.rs
index 3af740e819193a..c8f89421757c5c 100644
--- a/crates/ruff_linter/src/noqa.rs
+++ b/crates/ruff_linter/src/noqa.rs
@@ -664,6 +664,7 @@ impl Ranged for NoqaDirectiveLine<'_> {
 #[derive(Debug, Default)]
 pub(crate) struct NoqaDirectives<'a> {
     inner: Vec<NoqaDirectiveLine<'a>>,
+    last_directive_includes_eof: bool,
 }
 
 impl<'a> NoqaDirectives<'a> {
@@ -694,14 +695,17 @@ impl<'a> NoqaDirectives<'a> {
             }
         }
 
-        // Extend a mapping at the end of the file to also include the EOF token.
-        if let Some(last) = directives.last_mut() {
-            if last.range.end() == locator.contents().text_len() {
-                last.range = last.range.add_end(TextSize::from(1));
-            }
-        }
+        // Record whether last directive should include EOF token.
+        let last_directive_includes_eof = if let Some(last) = directives.last() {
+            last.range.end() == locator.contents().text_len()
+        } else {
+            false
+        };
 
-        Self { inner: directives }
+        Self {
+            inner: directives,
+            last_directive_includes_eof,
+        }
     }
 
     pub(crate) fn find_line_with_directive(&self, offset: TextSize) -> Option<&NoqaDirectiveLine> {
@@ -720,17 +724,25 @@ impl<'a> NoqaDirectives<'a> {
     }
 
     fn find_line_index(&self, offset: TextSize) -> Option<usize> {
-        self.inner
-            .binary_search_by(|directive| {
-                if directive.range.end() < offset {
-                    std::cmp::Ordering::Less
-                } else if directive.range.contains(offset) {
-                    std::cmp::Ordering::Equal
-                } else {
-                    std::cmp::Ordering::Greater
+        let index = self.inner.binary_search_by(|directive| {
+            if directive.range.end() < offset {
+                std::cmp::Ordering::Less
+            } else if directive.range.contains(offset) {
+                std::cmp::Ordering::Equal
+            } else {
+                std::cmp::Ordering::Greater
+            }
+        });
+
+        match index {
+            Ok(index) => Some(index),
+            Err(index) => {
+                if self.last_directive_includes_eof && index == self.inner.len() - 1 {
+                    return Some(index);
                 }
-            })
-            .ok()
+                None
+            }
+        }
     }
 
     pub(crate) fn lines(&self) -> &[NoqaDirectiveLine] {
diff --git a/crates/ruff_linter/src/rules/pygrep_hooks/rules/blanket_noqa.rs b/crates/ruff_linter/src/rules/pygrep_hooks/rules/blanket_noqa.rs
index 0213d4d6a920ca..d33560a7f90d5a 100644
--- a/crates/ruff_linter/src/rules/pygrep_hooks/rules/blanket_noqa.rs
+++ b/crates/ruff_linter/src/rules/pygrep_hooks/rules/blanket_noqa.rs
@@ -88,14 +88,8 @@ pub(crate) fn blanket_noqa(
 ) {
     for directive_line in noqa_directives.lines() {
         if let Directive::All(all) = &directive_line.directive {
-            let line = if directive_line.end() > locator.contents().text_len() {
-                let range = TextRange::new(directive_line.start(), locator.contents().text_len());
-                locator.slice(range)
-            } else {
-                locator.slice(directive_line)
-            };
-            let offset = directive_line.start();
-            let noqa_end = all.end() - offset;
+            let line = locator.slice(directive_line);
+            let noqa_end = all.end() - directive_line.start();
 
             // Skip the `# noqa`, plus any trailing whitespace.
             let mut cursor = Cursor::new(&line[noqa_end.to_usize()..]);