Support file-level type: ignore comments (#15081)

crates/red_knot_python_semantic/resources/mdtest/suppressions/knot-ignore.md

@@ -118,3 +118,15 @@ An empty codes array suppresses no-diagnostics and is always useless
 # error: [division-by-zero]
 a = 4 / 0  # knot: ignore[]
 ```
+
+## File-level suppression comments
+
+File level suppression comments are currently intentionally unsupported because we've yet to decide
+if they should use a different syntax that also supports enabling rules or changing the rule's
+severity: `knot: possibly-undefined-reference=error`
+
+```py
+# knot: ignore[division-by-zero]
+
+a = 4 / 0  # error: [division-by-zero]
+```

crates/red_knot_python_semantic/resources/mdtest/suppressions/type-ignore.md

@@ -122,3 +122,36 @@ a = (  # type: ignore
     test + 4  # error: [unresolved-reference]
 )
 ```
+
+## File level suppression
+
+```py
+# type: ignore
+
+a = 10 / 0
+b = a / 0
+```
+
+## File level suppression with leading shebang
+
+```py
+#!/usr/bin/env/python
+# type: ignore
+
+a = 10 / 0
+b = a / 0
+```
+
+## Invalid own-line suppression
+
+```py
+"""
+File level suppressions must come before any non-trivia token,
+including module docstrings. 
+"""
+
+# type: ignore
+
+a = 10 / 0  # error: [division-by-zero]
+b = a / 0  # error: [division-by-zero]
+```

crates/red_knot_python_semantic/src/suppression.rs

@@ -1,80 +1,31 @@
 use ruff_db::{files::File, parsed::parsed_module, source::source_text};
 use ruff_python_parser::TokenKind;
 use ruff_python_trivia::Cursor;
-use ruff_source_file::LineRanges;
+use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
-use ruff_text_size::{Ranged, TextRange, TextSize};
 use smallvec::{smallvec, SmallVec};
 
+use crate::lint::LintRegistry;
 use crate::{lint::LintId, Db};
 
 #[salsa::tracked(return_ref)]
 pub(crate) fn suppressions(db: &dyn Db, file: File) -> Suppressions {
-    let source = source_text(db.upcast(), file);
     let parsed = parsed_module(db.upcast(), file);
+    let source = source_text(db.upcast(), file);
 
-    let lints = db.lint_registry();
+    let mut builder = SuppressionsBuilder::new(&source, db.lint_registry());
-
+    let mut line_start = TextSize::default();
-    // TODO: Support `type: ignore` comments at the
-    //   [start of the file](https://typing.readthedocs.io/en/latest/spec/directives.html#type-ignore-comments).
-    let mut suppressions = Vec::default();
-    let mut line_start = source.bom_start_offset();
 
     for token in parsed.tokens() {
+        if !token.kind().is_trivia() {
+            builder.set_seen_non_trivia_token();
+        }
+
         match token.kind() {
             TokenKind::Comment => {
                 let parser = SuppressionParser::new(&source, token.range());
-                let suppressed_range = TextRange::new(line_start, token.range().end());
 
                 for comment in parser {
-                    match comment.codes {
+                    builder.add_comment(comment, line_start);
-                        // `type: ignore`
-                        None => {
-                            suppressions.push(Suppression {
-                                target: SuppressionTarget::All,
-                                comment_range: comment.range,
-                                range: comment.range,
-                                suppressed_range,
-                            });
-                        }
-
-                        // `type: ignore[..]`
-                        // The suppression applies to all lints if it is a `type: ignore`
-                        // comment. `type: ignore` apply to all lints for better mypy compatibility.
-                        Some(_) if comment.kind.is_type_ignore() => {
-                            suppressions.push(Suppression {
-                                target: SuppressionTarget::All,
-                                comment_range: comment.range,
-                                range: comment.range,
-                                suppressed_range,
-                            });
-                        }
-
-                        // `knot: ignore[a, b]`
-                        Some(codes) => {
-                            for code in &codes {
-                                match lints.get(&source[*code]) {
-                                    Ok(lint) => {
-                                        let range = if codes.len() == 1 {
-                                            comment.range
-                                        } else {
-                                            *code
-                                        };
-
-                                        suppressions.push(Suppression {
-                                            target: SuppressionTarget::Lint(lint),
-                                            range,
-                                            comment_range: comment.range,
-                                            suppressed_range,
-                                        });
-                                    }
-                                    Err(error) => {
-                                        tracing::debug!("Invalid suppression: {error}");
-                                        // TODO(micha): Handle invalid lint codes
-                                    }
-                                }
-                            }
-                        }
-                    }
                 }
             }
             TokenKind::Newline | TokenKind::NonLogicalNewline => {
@@ -84,34 +35,46 @@ pub(crate) fn suppressions(db: &dyn Db, file: File) -> Suppressions {
         }
     }
 
-    Suppressions { suppressions }
+    builder.finish()
 }
 
-/// The suppression of a single file.
+/// The suppressions of a single file.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub(crate) struct Suppressions {
-    /// The suppressions sorted by the suppressed range.
+    /// Suppressions that apply to the entire file.
     ///
-    /// It's possible that multiple suppressions apply for the same range.
+    /// The suppressions are sorted by [`Suppression::comment_range`] and the [`Suppression::suppressed_range`]
-    suppressions: Vec<Suppression>,
+    /// spans the entire file.
+    ///
+    /// For now, this is limited to `type: ignore` comments.
+    file: Vec<Suppression>,
+
+    /// Suppressions that apply to a specific line (or lines).
+    ///
+    /// Comments with multiple codes create multiple [`Suppression`]s that all share the same [`Suppression::comment_range`].
+    ///
+    /// The suppressions are sorted by [`Suppression::range`] (which implies [`Suppression::comment_range`]).
+    line: Vec<Suppression>,
 }
 
 impl Suppressions {
     pub(crate) fn find_suppression(&self, range: TextRange, id: LintId) -> Option<&Suppression> {
-        self.for_range(range)
+        self.file
+            .iter()
+            .chain(self.line_suppressions(range))
             .find(|suppression| suppression.matches(id))
     }
 
-    /// Returns all suppression comments that apply for `range`.
+    /// Returns the line-level suppressions that apply for `range`.
     ///
     /// A suppression applies for the given range if it contains the range's
     /// start or end offset. This means the suppression is on the same line
     /// as the diagnostic's start or end.
-    fn for_range(&self, range: TextRange) -> impl Iterator<Item = &Suppression> + '_ {
+    fn line_suppressions(&self, range: TextRange) -> impl Iterator<Item = &Suppression> + '_ {
         // First find the index of the suppression comment that ends right before the range
         // starts. This allows us to skip suppressions that are not relevant for the range.
         let end_offset = self
-            .suppressions
+            .line
             .binary_search_by_key(&range.start(), |suppression| {
                 suppression.suppressed_range.end()
             })
@@ -120,7 +83,7 @@ impl Suppressions {
         // From here, search the remaining suppression comments for one that
         // contains the range's start or end offset. Stop the search
         // as soon as the suppression's range and the range no longer overlap.
-        self.suppressions[end_offset..]
+        self.line[end_offset..]
             .iter()
             // Stop searching if the suppression starts after the range we're looking for.
             .take_while(move |suppression| range.end() >= suppression.suppressed_range.start())
@@ -177,6 +140,116 @@ enum SuppressionTarget {
     Lint(LintId),
 }
 
+struct SuppressionsBuilder<'a> {
+    lint_registry: &'a LintRegistry,
+    source: &'a str,
+
+    /// `type: ignore` comments at the top of the file before any non-trivia code apply to the entire file.
+    /// This boolean tracks if there has been any non trivia token.
+    seen_non_trivia_token: bool,
+
+    line: Vec<Suppression>,
+    file: Vec<Suppression>,
+}
+
+impl<'a> SuppressionsBuilder<'a> {
+    fn new(source: &'a str, lint_registry: &'a LintRegistry) -> Self {
+        Self {
+            source,
+            lint_registry,
+            seen_non_trivia_token: false,
+            line: Vec::new(),
+            file: Vec::new(),
+        }
+    }
+
+    fn set_seen_non_trivia_token(&mut self) {
+        self.seen_non_trivia_token = true;
+    }
+
+    fn finish(mut self) -> Suppressions {
+        self.line.shrink_to_fit();
+        self.file.shrink_to_fit();
+
+        Suppressions {
+            file: self.file,
+            line: self.line,
+        }
+    }
+
+    fn add_comment(&mut self, comment: SuppressionComment, line_start: TextSize) {
+        let (suppressions, suppressed_range) =
+            // `type: ignore` comments at the start of the file apply to the entire range.
+            // > A # type: ignore comment on a line by itself at the top of a file, before any docstrings,
+            // > imports, or other executable code, silences all errors in the file.
+            // > Blank lines and other comments, such as shebang lines and coding cookies,
+            // > may precede the # type: ignore comment.
+            // > https://typing.readthedocs.io/en/latest/spec/directives.html#type-ignore-comments
+            if comment.kind.is_type_ignore() && !self.seen_non_trivia_token {
+                (
+                    &mut self.file,
+                    TextRange::new(0.into(), self.source.text_len()),
+                )
+            } else {
+                (
+                    &mut self.line,
+                    TextRange::new(line_start, comment.range.end()),
+                )
+            };
+
+        match comment.codes {
+            // `type: ignore`
+            None => {
+                suppressions.push(Suppression {
+                    target: SuppressionTarget::All,
+                    comment_range: comment.range,
+                    range: comment.range,
+                    suppressed_range,
+                });
+            }
+
+            // `type: ignore[..]`
+            // The suppression applies to all lints if it is a `type: ignore`
+            // comment. `type: ignore` apply to all lints for better mypy compatibility.
+            Some(_) if comment.kind.is_type_ignore() => {
+                suppressions.push(Suppression {
+                    target: SuppressionTarget::All,
+                    comment_range: comment.range,
+                    range: comment.range,
+                    suppressed_range,
+                });
+            }
+
+            // `knot: ignore[a, b]`
+            Some(codes) => {
+                for code_range in &codes {
+                    let code = &self.source[*code_range];
+                    match self.lint_registry.get(code) {
+                        Ok(lint) => {
+                            let range = if codes.len() == 1 {
+                                comment.range
+                            } else {
+                                *code_range
+                            };
+
+                            suppressions.push(Suppression {
+                                target: SuppressionTarget::Lint(lint),
+                                range,
+                                comment_range: comment.range,
+                                suppressed_range,
+                            });
+                        }
+                        Err(error) => {
+                            tracing::debug!("Invalid suppression: {error}");
+                            // TODO(micha): Handle invalid lint codes
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
 struct SuppressionParser<'src> {
     cursor: Cursor<'src>,
     range: TextRange,