From d3b663c932886e39cf44e455e9237ec0f6cfb7c2 Mon Sep 17 00:00:00 2001 From: dylwil3 Date: Mon, 17 Nov 2025 17:30:24 -0600 Subject: [PATCH] see if text ranges work instead --- .../ruff_python_formatter/src/comments/mod.rs | 49 +++++++++++++++---- crates/ruff_python_formatter/src/context.rs | 6 +-- crates/ruff_python_formatter/src/lib.rs | 5 +- crates/ruff_python_formatter/src/range.rs | 4 +- .../src/statement/suite.rs | 4 +- .../src/string/docstring.rs | 4 +- 6 files changed, 51 insertions(+), 21 deletions(-) diff --git a/crates/ruff_python_formatter/src/comments/mod.rs b/crates/ruff_python_formatter/src/comments/mod.rs index 6c93edd35c..cb399bd8b3 100644 --- a/crates/ruff_python_formatter/src/comments/mod.rs +++ b/crates/ruff_python_formatter/src/comments/mod.rs @@ -99,7 +99,6 @@ 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}; @@ -512,21 +511,51 @@ pub(crate) fn has_skip_comment(trailing_comments: &[SourceComment], source: &str }) } -pub(crate) struct SuppressedNodes<'a>(FxHashSet>); +pub(crate) struct SuppressedNodeRanges(Vec); -impl<'a> SuppressedNodes<'a> { +impl<'a> SuppressedNodeRanges { pub(crate) fn from_comments(comments: &Comments<'a>, source: &'a str) -> Self { let map = &comments.clone().data.comments; - Self( - map.keys() - .copied() - .filter(|key| has_skip_comment(map.trailing(key), source)) - .collect(), - ) + + let mut ranges = map + .keys() + .copied() + .filter_map(|key| suppressed_range(key.node(), map.trailing(&key), source)) + .collect::>(); + + ranges.sort_by_key(ruff_text_size::Ranged::start); + + Self(ranges) } pub(crate) fn contains(&self, node: AnyNodeRef<'_>) -> bool { - self.0.contains(&NodeRefEqualityKey::from_ref(node)) + 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 { + if has_skip_comment(trailing_comments, source) { + let end = node.end(); + let start = node.start(); + + Some(TextRange::new(start, end)) + } else { + None } } diff --git a/crates/ruff_python_formatter/src/context.rs b/crates/ruff_python_formatter/src/context.rs index 19164bbff4..35b7f3132a 100644 --- a/crates/ruff_python_formatter/src/context.rs +++ b/crates/ruff_python_formatter/src/context.rs @@ -7,14 +7,14 @@ use ruff_python_ast::str::Quote; use ruff_python_parser::Tokens; use crate::PyFormatOptions; -use crate::comments::{Comments, SuppressedNodes}; +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: SuppressedNodes<'a>, + suppressed_nodes: SuppressedNodeRanges, tokens: &'a Tokens, node_level: NodeLevel, indent_level: IndentLevel, @@ -36,7 +36,7 @@ impl<'a> PyFormatContext<'a> { options: PyFormatOptions, contents: &'a str, comments: Comments<'a>, - suppressed_nodes: SuppressedNodes<'a>, + suppressed_nodes: SuppressedNodeRanges, tokens: &'a Tokens, ) -> Self { Self { diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs index b9206ff9a5..3770bcf72b 100644 --- a/crates/ruff_python_formatter/src/lib.rs +++ b/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, SuppressedNodes, 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; @@ -174,7 +175,7 @@ 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 suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source); let formatted = format!( PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens()), diff --git a/crates/ruff_python_formatter/src/range.rs b/crates/ruff_python_formatter/src/range.rs index 11d1dc6fba..f77529bf05 100644 --- a/crates/ruff_python_formatter/src/range.rs +++ b/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, SuppressedNodes}; +use crate::comments::{Comments, SuppressedNodeRanges}; use crate::context::{IndentLevel, NodeLevel}; use crate::prelude::*; use crate::statement::suite::DocstringStmt; @@ -77,7 +77,7 @@ 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 suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source); let mut context = PyFormatContext::new( options.with_source_map_generation(SourceMapGeneration::Enabled), diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs index e4399280ae..999983f230 100644 --- a/crates/ruff_python_formatter/src/statement/suite.rs +++ b/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, SuppressedNodes}; + use crate::comments::{Comments, SuppressedNodeRanges}; use crate::prelude::*; use crate::statement::suite::SuiteKind; @@ -974,7 +974,7 @@ 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 suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source); let context = PyFormatContext::new( PyFormatOptions::default(), diff --git a/crates/ruff_python_formatter/src/string/docstring.rs b/crates/ruff_python_formatter/src/string/docstring.rs index f642519c86..f3d98ae1bd 100644 --- a/crates/ruff_python_formatter/src/string/docstring.rs +++ b/crates/ruff_python_formatter/src/string/docstring.rs @@ -20,7 +20,7 @@ use { }; use super::NormalizedString; -use crate::comments::SuppressedNodes; +use crate::comments::SuppressedNodeRanges; use crate::preview::is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled; use crate::string::StringQuotes; use crate::{DocstringCodeLineWidth, FormatModuleError, prelude::*}; @@ -1586,7 +1586,7 @@ 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 suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source_code.as_str()); let ctx = PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens()) .in_docstring(docstring_quote_style);