[WIP] Respect fmt: skip for multiple statements on a single line

crates/ruff_python_formatter/src/comments/mod.rs

@@ -511,6 +511,54 @@ pub(crate) fn has_skip_comment(trailing_comments: &[SourceComment], source: &str
     })
 }
 
+pub(crate) struct SuppressedNodeRanges(Vec<TextRange>);
+
+impl<'a> SuppressedNodeRanges {
+    pub(crate) fn from_comments(comments: &Comments<'a>, source: &'a str) -> Self {
+        let map = &comments.clone().data.comments;
+
+        let mut ranges = map
+            .keys()
+            .copied()
+            .filter_map(|key| suppressed_range(key.node(), map.trailing(&key), source))
+            .collect::<Vec<_>>();
+
+        ranges.sort_by_key(ruff_text_size::Ranged::start);
+
+        Self(ranges)
+    }
+
+    pub(crate) fn contains(&self, node: AnyNodeRef<'_>) -> bool {
+        let target = node.range();
+        self.0
+            .binary_search_by(|range| {
+                if range.eq(&target) {
+                    std::cmp::Ordering::Equal
+                } else if range.start() < target.start() {
+                    std::cmp::Ordering::Less
+                } else {
+                    std::cmp::Ordering::Greater
+                }
+            })
+            .is_ok()
+    }
+}
+
+fn suppressed_range<'a>(
+    node: AnyNodeRef<'a>,
+    trailing_comments: &[SourceComment],
+    source: &'a str,
+) -> Option<TextRange> {
+    if has_skip_comment(trailing_comments, source) {
+        let end = node.end();
+        let start = node.start();
+
+        Some(TextRange::new(start, end))
+    } else {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use insta::assert_debug_snapshot;

crates/ruff_python_formatter/src/context.rs

@@ -2,17 +2,19 @@ use std::fmt::{Debug, Formatter};
 use std::ops::{Deref, DerefMut};
 
 use ruff_formatter::{Buffer, FormatContext, GroupId, IndentWidth, SourceCode};
+use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::str::Quote;
 use ruff_python_parser::Tokens;
 
 use crate::PyFormatOptions;
-use crate::comments::Comments;
+use crate::comments::{Comments, SuppressedNodeRanges};
 use crate::other::interpolated_string::InterpolatedStringContext;
 
 pub struct PyFormatContext<'a> {
     options: PyFormatOptions,
     contents: &'a str,
     comments: Comments<'a>,
+    suppressed_nodes: SuppressedNodeRanges,
     tokens: &'a Tokens,
     node_level: NodeLevel,
     indent_level: IndentLevel,
@@ -34,12 +36,14 @@ impl<'a> PyFormatContext<'a> {
         options: PyFormatOptions,
         contents: &'a str,
         comments: Comments<'a>,
+        suppressed_nodes: SuppressedNodeRanges,
         tokens: &'a Tokens,
     ) -> Self {
         Self {
             options,
             contents,
             comments,
+            suppressed_nodes,
             tokens,
             node_level: NodeLevel::TopLevel(TopLevelStatementPosition::Other),
             indent_level: IndentLevel::new(0),
@@ -112,6 +116,11 @@ impl<'a> PyFormatContext<'a> {
     pub(crate) const fn is_preview(&self) -> bool {
         self.options.preview().is_enabled()
     }
+
+    /// Returns `true` if node is suppressed via skip comment.
+    pub(crate) fn is_suppressed(&self, node: AnyNodeRef<'_>) -> bool {
+        self.suppressed_nodes.contains(node)
+    }
 }
 
 impl FormatContext for PyFormatContext<'_> {

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,8 @@ 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, SuppressedNodeRanges, has_skip_comment, leading_comments,
+    trailing_comments,
 };
 pub use crate::context::PyFormatContext;
 pub use crate::db::Db;
@@ -61,7 +62,7 @@ where
         let node_ref = AnyNodeRef::from(node);
         let node_comments = comments.leading_dangling_trailing(node_ref);
 
-        if self.is_suppressed(node_comments.trailing, f.context()) {
+        if self.is_suppressed(node, f.context()) {
             suppressed_node(node_ref).fmt(f)
         } else {
             leading_comments(node_comments.leading).fmt(f)?;
@@ -99,11 +100,7 @@ where
     /// Formats the node's fields.
     fn fmt_fields(&self, item: &N, f: &mut PyFormatter) -> FormatResult<()>;
 
-    fn is_suppressed(
-        &self,
-        _trailing_comments: &[SourceComment],
-        _context: &PyFormatContext,
-    ) -> bool {
+    fn is_suppressed(&self, _node: &N, _context: &PyFormatContext) -> bool {
         false
     }
 }
@@ -178,9 +175,10 @@ where
 {
     let source_code = SourceCode::new(source);
     let comments = Comments::from_ast(parsed.syntax(), source_code, comment_ranges);
+    let suppressed_nodes = SuppressedNodeRanges::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/other/decorator.rs

@@ -1,10 +1,9 @@
 use ruff_formatter::write;
 use ruff_python_ast::Decorator;
 
-use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
-use crate::{has_skip_comment, prelude::*};
+use crate::prelude::*;
 
 #[derive(Default)]
 pub struct FormatDecorator;
@@ -25,12 +24,7 @@ impl FormatNodeRule<Decorator> for FormatDecorator {
             ]
         )
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &Decorator, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

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, SuppressedNodeRanges};
 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 = SuppressedNodeRanges::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/stmt_ann_assign.rs

@@ -1,14 +1,13 @@
 use ruff_formatter::write;
 use ruff_python_ast::StmtAnnAssign;
 
-use crate::comments::SourceComment;
 use crate::expression::is_splittable_expression;
 use crate::expression::parentheses::Parentheses;
+use crate::prelude::*;
 use crate::statement::stmt_assign::{
     AnyAssignmentOperator, AnyBeforeOperator, FormatStatementsLastExpression,
 };
 use crate::statement::trailing_semicolon;
-use crate::{has_skip_comment, prelude::*};
 
 #[derive(Default)]
 pub struct FormatStmtAnnAssign;
@@ -84,12 +83,7 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
 
         Ok(())
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtAnnAssign, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_assert.rs

@@ -2,10 +2,9 @@ use ruff_formatter::prelude::{space, token};
 use ruff_formatter::write;
 use ruff_python_ast::StmtAssert;
 
-use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
-use crate::{has_skip_comment, prelude::*};
+use crate::prelude::*;
 
 #[derive(Default)]
 pub struct FormatStmtAssert;
@@ -41,12 +40,7 @@ impl FormatNodeRule<StmtAssert> for FormatStmtAssert {
 
         Ok(())
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtAssert, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -18,13 +18,13 @@ use crate::expression::{
     maybe_parenthesize_expression,
 };
 use crate::other::interpolated_string::InterpolatedStringLayout;
+use crate::prelude::*;
 use crate::statement::trailing_semicolon;
 use crate::string::StringLikeExtensions;
 use crate::string::implicit::{
     FormatImplicitConcatenatedStringExpanded, FormatImplicitConcatenatedStringFlat,
     ImplicitConcatenatedLayout,
 };
-use crate::{has_skip_comment, prelude::*};
 
 #[derive(Default)]
 pub struct FormatStmtAssign;
@@ -102,13 +102,8 @@ impl FormatNodeRule<StmtAssign> for FormatStmtAssign {
 
         Ok(())
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtAssign, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }
 

crates/ruff_python_formatter/src/statement/stmt_aug_assign.rs

@@ -1,15 +1,14 @@
 use ruff_formatter::write;
 use ruff_python_ast::StmtAugAssign;
 
-use crate::comments::SourceComment;
 use crate::expression::parentheses::is_expression_parenthesized;
+use crate::prelude::*;
 use crate::statement::stmt_assign::{
     AnyAssignmentOperator, AnyBeforeOperator, FormatStatementsLastExpression,
     has_target_own_parentheses,
 };
 use crate::statement::trailing_semicolon;
 use crate::{AsFormat, FormatNodeRule};
-use crate::{has_skip_comment, prelude::*};
 
 #[derive(Default)]
 pub struct FormatStmtAugAssign;
@@ -62,12 +61,7 @@ impl FormatNodeRule<StmtAugAssign> for FormatStmtAugAssign {
 
         Ok(())
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtAugAssign, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_break.rs

@@ -1,7 +1,6 @@
 use ruff_python_ast::StmtBreak;
 
-use crate::comments::SourceComment;
-use crate::{has_skip_comment, prelude::*};
+use crate::prelude::*;
 
 #[derive(Default)]
 pub struct FormatStmtBreak;
@@ -10,12 +9,7 @@ impl FormatNodeRule<StmtBreak> for FormatStmtBreak {
     fn fmt_fields(&self, _item: &StmtBreak, f: &mut PyFormatter) -> FormatResult<()> {
         token("break").fmt(f)
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtBreak, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_continue.rs

@@ -1,7 +1,6 @@
 use ruff_python_ast::StmtContinue;
 
-use crate::comments::SourceComment;
-use crate::{has_skip_comment, prelude::*};
+use crate::prelude::*;
 
 #[derive(Default)]
 pub struct FormatStmtContinue;
@@ -10,12 +9,7 @@ impl FormatNodeRule<StmtContinue> for FormatStmtContinue {
     fn fmt_fields(&self, _item: &StmtContinue, f: &mut PyFormatter) -> FormatResult<()> {
         token("continue").fmt(f)
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtContinue, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_delete.rs

@@ -3,10 +3,10 @@ use ruff_python_ast::StmtDelete;
 use ruff_text_size::Ranged;
 
 use crate::builders::{PyFormatterExtensions, parenthesize_if_expands};
-use crate::comments::{SourceComment, dangling_node_comments};
+use crate::comments::dangling_node_comments;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
-use crate::{has_skip_comment, prelude::*};
+use crate::prelude::*;
 
 #[derive(Default)]
 pub struct FormatStmtDelete;
@@ -57,12 +57,7 @@ impl FormatNodeRule<StmtDelete> for FormatStmtDelete {
             }
         }
     }
-
-    fn is_suppressed(
-        &self,
-        trailing_comments: &[SourceComment],
-        context: &PyFormatContext,
-    ) -> bool {
-        has_skip_comment(trailing_comments, context.source())
+    fn is_suppressed(&self, node: &StmtDelete, context: &PyFormatContext) -> bool {
+        context.is_suppressed(node.into())
     }
 }