Add token based parenthesized_ranges implementation (#21738)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_python_ast/src/parenthesize.rs

@@ -11,6 +11,8 @@ use crate::ExprRef;
 /// Note that without a parent the range can be inaccurate, e.g. `f(a)` we falsely return a set of
 /// parentheses around `a` even if the parentheses actually belong to `f`. That is why you should
 /// generally prefer [`parenthesized_range`].
+///
+/// Prefer [`crate::token::parentheses_iterator`] if you have access to [`crate::token::Tokens`].
 pub fn parentheses_iterator<'a>(
     expr: ExprRef<'a>,
     parent: Option<AnyNodeRef>,
@@ -57,6 +59,8 @@ pub fn parentheses_iterator<'a>(
 
 /// Returns the [`TextRange`] of a given expression including parentheses, if the expression is
 /// parenthesized; or `None`, if the expression is not parenthesized.
+///
+/// Prefer [`crate::token::parenthesized_range`] if you have access to [`crate::token::Tokens`].
 pub fn parenthesized_range(
     expr: ExprRef,
     parent: AnyNodeRef,

crates/ruff_python_ast/src/token.rs

@@ -16,8 +16,10 @@ use crate::str_prefix::{
 use crate::{AnyStringFlags, BoolOp, Operator, StringFlags, UnaryOp};
 use ruff_text_size::{Ranged, TextRange};
 
+mod parentheses;
 mod tokens;
 
+pub use parentheses::{parentheses_iterator, parenthesized_range};
 pub use tokens::{TokenAt, TokenIterWithContext, Tokens};
 
 #[derive(Clone, Copy, PartialEq, Eq)]

crates/ruff_python_ast/src/token/parentheses.rs

@@ -0,0 +1,58 @@
+use ruff_text_size::{Ranged, TextLen, TextRange};
+
+use super::{TokenKind, Tokens};
+use crate::{AnyNodeRef, ExprRef};
+
+/// Returns an iterator over the ranges of the optional parentheses surrounding an expression.
+///
+/// E.g. for `((f()))` with `f()` as expression, the iterator returns the ranges (1, 6) and (0, 7).
+///
+/// Note that without a parent the range can be inaccurate, e.g. `f(a)` we falsely return a set of
+/// parentheses around `a` even if the parentheses actually belong to `f`. That is why you should
+/// generally prefer [`parenthesized_range`].
+pub fn parentheses_iterator<'a>(
+    expr: ExprRef<'a>,
+    parent: Option<AnyNodeRef>,
+    tokens: &'a Tokens,
+) -> impl Iterator<Item = TextRange> + 'a {
+    let after_tokens = if let Some(parent) = parent {
+        // If the parent is a node that brings its own parentheses, exclude the closing parenthesis
+        // from our search range. Otherwise, we risk matching on calls, like `func(x)`, for which
+        // the open and close parentheses are part of the `Arguments` node.
+        let exclusive_parent_end = if parent.is_arguments() {
+            parent.end() - ")".text_len()
+        } else {
+            parent.end()
+        };
+
+        tokens.in_range(TextRange::new(expr.end(), exclusive_parent_end))
+    } else {
+        tokens.after(expr.end())
+    };
+
+    let right_parens = after_tokens
+        .iter()
+        .filter(|token| !token.kind().is_trivia())
+        .take_while(move |token| token.kind() == TokenKind::Rpar);
+
+    let left_parens = tokens
+        .before(expr.start())
+        .iter()
+        .rev()
+        .filter(|token| !token.kind().is_trivia())
+        .take_while(|token| token.kind() == TokenKind::Lpar);
+
+    right_parens
+        .zip(left_parens)
+        .map(|(right, left)| TextRange::new(left.start(), right.end()))
+}
+
+/// Returns the [`TextRange`] of a given expression including parentheses, if the expression is
+/// parenthesized; or `None`, if the expression is not parenthesized.
+pub fn parenthesized_range(
+    expr: ExprRef,
+    parent: AnyNodeRef,
+    tokens: &Tokens,
+) -> Option<TextRange> {
+    parentheses_iterator(expr, Some(parent), tokens).last()
+}