store suppressed nodes ahead of time

Author: dylwil3

crates/ruff_python_formatter/src/comments/mod.rs

@@ -99,6 +99,7 @@ use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_trivia::{CommentLinePosition, CommentRanges, SuppressionKind};
 use ruff_text_size::{Ranged, TextRange};
+use rustc_hash::FxHashSet;
 pub(crate) use visitor::collect_comments;
 
 use crate::comments::debug::{DebugComment, DebugComments};
@@ -511,6 +512,23 @@ pub(crate) fn has_skip_comment(trailing_comments: &[SourceComment], source: &str
     })
 }
 
+pub(crate) struct SuppressedNodes<'a>(FxHashSet<NodeRefEqualityKey<'a>>);
+
+impl<'a> SuppressedNodes<'a> {
+    pub(crate) fn from_comments(comments: &Comments<'a>, source: &'a str) -> Self {
+        let map = &comments.clone().data.comments;
+        Self(FxHashSet::from_iter(
+            map.keys()
+                .copied()
+                .filter(|key| has_skip_comment(map.trailing(key), source)),
+        ))
+    }
+
+    pub(crate) fn contains(&self, node: AnyNodeRef<'_>) -> bool {
+        self.0.contains(&NodeRefEqualityKey::from_ref(node))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use insta::assert_debug_snapshot;

crates/ruff_python_formatter/src/context.rs

@@ -7,13 +7,14 @@ use ruff_python_ast::str::Quote;
 use ruff_python_parser::Tokens;
 
 use crate::PyFormatOptions;
-use crate::comments::{Comments, has_skip_comment};
+use crate::comments::{Comments, SuppressedNodes};
 use crate::other::interpolated_string::InterpolatedStringContext;
 
 pub struct PyFormatContext<'a> {
     options: PyFormatOptions,
     contents: &'a str,
     comments: Comments<'a>,
+    suppressed_nodes: SuppressedNodes<'a>,
     tokens: &'a Tokens,
     node_level: NodeLevel,
     indent_level: IndentLevel,
@@ -35,12 +36,14 @@ impl<'a> PyFormatContext<'a> {
         options: PyFormatOptions,
         contents: &'a str,
         comments: Comments<'a>,
+        suppressed_nodes: SuppressedNodes<'a>,
         tokens: &'a Tokens,
     ) -> Self {
         Self {
             options,
             contents,
             comments,
+            suppressed_nodes,
             tokens,
             node_level: NodeLevel::TopLevel(TopLevelStatementPosition::Other),
             indent_level: IndentLevel::new(0),
@@ -116,7 +119,7 @@ impl<'a> PyFormatContext<'a> {
 
     /// Returns `true` if node is suppressed via skip comment.
     pub(crate) fn is_suppressed(&self, node: AnyNodeRef<'_>) -> bool {
-        has_skip_comment(self.comments().trailing(node), self.source())
+        self.suppressed_nodes.contains(node)
     }
 }
 

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,7 @@ use ruff_python_trivia::CommentRanges;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::comments::{
-    Comments, SourceComment, has_skip_comment, leading_comments, trailing_comments,
+    Comments, SourceComment, SuppressedNodes, has_skip_comment, leading_comments, trailing_comments,
 };
 pub use crate::context::PyFormatContext;
 pub use crate::db::Db;
@@ -174,9 +174,10 @@ where
 {
     let source_code = SourceCode::new(source);
     let comments = Comments::from_ast(parsed.syntax(), source_code, comment_ranges);
+    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
 
     let formatted = format!(
-        PyFormatContext::new(options, source, comments, parsed.tokens()),
+        PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens()),
         [parsed.syntax().format()]
     )?;
     formatted

crates/ruff_python_formatter/src/range.rs

@@ -12,7 +12,7 @@ use ruff_python_trivia::{
 };
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
-use crate::comments::Comments;
+use crate::comments::{Comments, SuppressedNodes};
 use crate::context::{IndentLevel, NodeLevel};
 use crate::prelude::*;
 use crate::statement::suite::DocstringStmt;
@@ -77,11 +77,13 @@ pub fn format_range(
     let source_code = SourceCode::new(source);
     let comment_ranges = CommentRanges::from(parsed.tokens());
     let comments = Comments::from_ast(parsed.syntax(), source_code, &comment_ranges);
+    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
 
     let mut context = PyFormatContext::new(
         options.with_source_map_generation(SourceMapGeneration::Enabled),
         source,
         comments,
+        suppressed_nodes,
         parsed.tokens(),
     );
 

crates/ruff_python_formatter/src/statement/suite.rs

@@ -945,7 +945,7 @@ mod tests {
     use ruff_python_trivia::CommentRanges;
 
     use crate::PyFormatOptions;
-    use crate::comments::Comments;
+    use crate::comments::{Comments, SuppressedNodes};
     use crate::prelude::*;
     use crate::statement::suite::SuiteKind;
 
@@ -973,11 +973,14 @@ def trailing_func():
 
         let parsed = parse_module(source).unwrap();
         let comment_ranges = CommentRanges::from(parsed.tokens());
+        let comments = Comments::from_ranges(&comment_ranges);
+        let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
 
         let context = PyFormatContext::new(
             PyFormatOptions::default(),
             source,
-            Comments::from_ranges(&comment_ranges),
+            comments,
+            suppressed_nodes,
             parsed.tokens(),
         );
 

crates/ruff_python_formatter/src/string/docstring.rs

@@ -20,6 +20,7 @@ use {
 };
 
 use super::NormalizedString;
+use crate::comments::SuppressedNodes;
 use crate::preview::is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled;
 use crate::string::StringQuotes;
 use crate::{DocstringCodeLineWidth, FormatModuleError, prelude::*};
@@ -1585,8 +1586,9 @@ fn docstring_format_source(
     let comment_ranges = CommentRanges::from(parsed.tokens());
     let source_code = ruff_formatter::SourceCode::new(source);
     let comments = crate::Comments::from_ast(parsed.syntax(), source_code, &comment_ranges);
+    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source_code.as_str());
 
-    let ctx = PyFormatContext::new(options, source, comments, parsed.tokens())
+    let ctx = PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens())
         .in_docstring(docstring_quote_style);
     let formatted = crate::format!(ctx, [parsed.syntax().format()])?;
     formatted