Make BoolOp its own located token (#3265)

crates/ruff_python_formatter/Cargo.toml

@@ -14,6 +14,7 @@ ruff_text_size = { path = "../ruff_text_size" }
 anyhow = { workspace = true }
 clap = { workspace = true }
 is-macro = { workspace = true }
+itertools = { workspace = true }
 once_cell = { workspace = true }
 rustc-hash = { workspace = true }
 rustpython-common = { workspace = true }

crates/ruff_python_formatter/src/attachment.rs

@@ -1,6 +1,8 @@
 use crate::core::visitor;
 use crate::core::visitor::Visitor;
-use crate::cst::{Alias, Arg, Body, Excepthandler, Expr, Pattern, SliceIndex, Stmt};
+use crate::cst::{
+    Alias, Arg, Body, BoolOp, Excepthandler, Expr, Keyword, Pattern, SliceIndex, Stmt,
+};
 use crate::trivia::{decorate_trivia, TriviaIndex, TriviaToken};
 
 struct AttachmentVisitor {
@@ -56,6 +58,22 @@ impl<'a> Visitor<'a> for AttachmentVisitor {
         visitor::walk_excepthandler(self, excepthandler);
     }
 
+    fn visit_keyword(&mut self, keyword: &'a mut Keyword) {
+        let trivia = self.index.keyword.remove(&keyword.id());
+        if let Some(comments) = trivia {
+            keyword.trivia.extend(comments);
+        }
+        visitor::walk_keyword(self, keyword);
+    }
+
+    fn visit_bool_op(&mut self, bool_op: &'a mut BoolOp) {
+        let trivia = self.index.bool_op.remove(&bool_op.id());
+        if let Some(comments) = trivia {
+            bool_op.trivia.extend(comments);
+        }
+        visitor::walk_bool_op(self, bool_op);
+    }
+
     fn visit_slice_index(&mut self, slice_index: &'a mut SliceIndex) {
         let trivia = self.index.slice_index.remove(&slice_index.id());
         if let Some(comments) = trivia {

crates/ruff_python_formatter/src/core/helpers.rs

@@ -28,6 +28,7 @@ pub fn trailing_quote(content: &str) -> Option<&&str> {
         .find(|&pattern| content.ends_with(pattern))
 }
 
+/// Return `true` if the given string is a radix literal (e.g., `0b101`).
 pub fn is_radix_literal(content: &str) -> bool {
     content.starts_with("0b")
         || content.starts_with("0o")
@@ -37,6 +38,29 @@ pub fn is_radix_literal(content: &str) -> bool {
         || content.starts_with("0X")
 }
 
+/// Find the first token in the given range that satisfies the given predicate.
+pub fn find_tok(
+    location: Location,
+    end_location: Location,
+    locator: &Locator,
+    f: impl Fn(rustpython_parser::Tok) -> bool,
+) -> (Location, Location) {
+    let (source, start_index, end_index) = locator.slice(Range::new(location, end_location));
+    for (start, tok, end) in rustpython_parser::lexer::lex_located(
+        &source[start_index..end_index],
+        rustpython_parser::Mode::Module,
+        location,
+    )
+    .flatten()
+    {
+        if f(tok) {
+            return (start, end);
+        }
+    }
+
+    unreachable!()
+}
+
 /// Expand the range of a compound statement.
 ///
 /// `location` is the start of the compound statement (e.g., the `if` in `if x:`).

crates/ruff_python_formatter/src/core/visitor.rs

@@ -1,7 +1,7 @@
 use rustpython_parser::ast::Constant;
 
 use crate::cst::{
-    Alias, Arg, Arguments, Body, Boolop, Cmpop, Comprehension, Excepthandler, ExcepthandlerKind,
+    Alias, Arg, Arguments, Body, BoolOp, Cmpop, Comprehension, Excepthandler, ExcepthandlerKind,
     Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern, PatternKind, SliceIndex,
     SliceIndexKind, Stmt, StmtKind, Unaryop, Withitem,
 };
@@ -22,8 +22,8 @@ pub trait Visitor<'a> {
     fn visit_expr_context(&mut self, expr_context: &'a mut ExprContext) {
         walk_expr_context(self, expr_context);
     }
-    fn visit_boolop(&mut self, boolop: &'a mut Boolop) {
-        walk_boolop(self, boolop);
+    fn visit_bool_op(&mut self, bool_op: &'a mut BoolOp) {
+        walk_bool_op(self, bool_op);
     }
     fn visit_operator(&mut self, operator: &'a mut Operator) {
         walk_operator(self, operator);
@@ -294,10 +294,12 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a mut Stm
 
 pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a mut Expr) {
     match &mut expr.node {
-        ExprKind::BoolOp { op, values } => {
-            visitor.visit_boolop(op);
-            for expr in values {
-                visitor.visit_expr(expr);
+        ExprKind::BoolOp { ops, values } => {
+            for op in ops {
+                visitor.visit_bool_op(op);
+            }
+            for value in values {
+                visitor.visit_expr(value);
             }
         }
         ExprKind::NamedExpr { target, value } => {
@@ -600,7 +602,7 @@ pub fn walk_expr_context<'a, V: Visitor<'a> + ?Sized>(
 }
 
 #[allow(unused_variables)]
-pub fn walk_boolop<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, boolop: &'a mut Boolop) {}
+pub fn walk_bool_op<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, bool_op: &'a mut BoolOp) {}
 
 #[allow(unused_variables)]
 pub fn walk_operator<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, operator: &'a mut Operator) {}

crates/ruff_python_formatter/src/cst.rs

@@ -1,9 +1,11 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 
-use crate::core::helpers::{expand_indented_block, is_elif};
 use rustpython_parser::ast::{Constant, Location};
 use rustpython_parser::Mode;
 
+use itertools::Itertools;
+
+use crate::core::helpers::{expand_indented_block, find_tok, is_elif};
 use crate::core::locator::Locator;
 use crate::core::types::Range;
 use crate::trivia::{Parenthesize, Trivia};
@@ -57,20 +59,22 @@ impl From<rustpython_parser::ast::ExprContext> for ExprContext {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum Boolop {
+pub enum BoolOpKind {
     And,
     Or,
 }
 
-impl From<rustpython_parser::ast::Boolop> for Boolop {
-    fn from(op: rustpython_parser::ast::Boolop) -> Self {
+impl From<&rustpython_parser::ast::Boolop> for BoolOpKind {
+    fn from(op: &rustpython_parser::ast::Boolop) -> Self {
         match op {
             rustpython_parser::ast::Boolop::And => Self::And,
             rustpython_parser::ast::Boolop::Or => Self::Or,
         }
     }
 }
 
+pub type BoolOp = Located<BoolOpKind>;
+
 #[derive(Clone, Debug, PartialEq)]
 pub enum Operator {
     Add,
@@ -308,7 +312,7 @@ pub type Stmt = Located<StmtKind>;
 #[derive(Clone, Debug, PartialEq)]
 pub enum ExprKind {
     BoolOp {
-        op: Boolop,
+        ops: Vec<BoolOp>,
         values: Vec<Expr>,
     },
     NamedExpr {
@@ -1677,7 +1681,23 @@ impl From<(rustpython_parser::ast::Expr, &Locator<'_>)> for Expr {
                 location: expr.location,
                 end_location: expr.end_location,
                 node: ExprKind::BoolOp {
-                    op: op.into(),
+                    ops: values
+                        .iter()
+                        .tuple_windows()
+                        .map(|(left, right)| {
+                            let target = match &op {
+                                rustpython_parser::ast::Boolop::And => rustpython_parser::Tok::And,
+                                rustpython_parser::ast::Boolop::Or => rustpython_parser::Tok::Or,
+                            };
+                            let (op_location, op_end_location) = find_tok(
+                                left.end_location.unwrap(),
+                                right.location,
+                                locator,
+                                |tok| tok == target,
+                            );
+                            BoolOp::new(op_location, op_end_location, (&op).into())
+                        })
+                        .collect(),
                     values: values
                         .into_iter()
                         .map(|node| (node, locator).into())

crates/ruff_python_formatter/src/format/bool_op.rs

@@ -0,0 +1,38 @@
+use ruff_formatter::prelude::*;
+use ruff_formatter::write;
+
+use crate::context::ASTFormatContext;
+use crate::cst::{BoolOp, BoolOpKind};
+use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
+use crate::shared_traits::AsFormat;
+
+pub struct FormatBoolOp<'a> {
+    item: &'a BoolOp,
+}
+
+impl AsFormat<ASTFormatContext<'_>> for BoolOp {
+    type Format<'a> = FormatBoolOp<'a>;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatBoolOp { item: self }
+    }
+}
+
+impl Format<ASTFormatContext<'_>> for FormatBoolOp<'_> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext>) -> FormatResult<()> {
+        let boolop = self.item;
+
+        write!(f, [leading_comments(boolop)])?;
+        write!(
+            f,
+            [text(match boolop.node {
+                BoolOpKind::And => "and",
+                BoolOpKind::Or => "or",
+            })]
+        )?;
+        write!(f, [end_of_line_comments(boolop)])?;
+        write!(f, [trailing_comments(boolop)])?;
+
+        Ok(())
+    }
+}

crates/ruff_python_formatter/src/format/expr.rs

@@ -9,7 +9,7 @@ use ruff_text_size::TextSize;
 use crate::context::ASTFormatContext;
 use crate::core::types::Range;
 use crate::cst::{
-    Arguments, Boolop, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, SliceIndex,
+    Arguments, BoolOp, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, SliceIndex,
     SliceIndexKind, Unaryop,
 };
 use crate::format::builders::literal;
@@ -338,38 +338,10 @@ fn format_call(
     if args.is_empty() && keywords.is_empty() {
         write!(f, [text("(")])?;
         write!(f, [text(")")])?;
-
-        // Format any end-of-line comments.
-        let mut first = true;
-        for range in expr.trivia.iter().filter_map(|trivia| {
-            if trivia.relationship.is_trailing() {
-                trivia.kind.end_of_line_comment()
-            } else {
-                None
-            }
-        }) {
-            if std::mem::take(&mut first) {
-                write!(f, [line_suffix(&text("  "))])?;
-            }
-            write!(f, [line_suffix(&literal(range))])?;
-        }
+        write!(f, [end_of_line_comments(expr)])?;
     } else {
         write!(f, [text("(")])?;
-
-        // Format any end-of-line comments.
-        let mut first = true;
-        for range in expr.trivia.iter().filter_map(|trivia| {
-            if trivia.relationship.is_trailing() {
-                trivia.kind.end_of_line_comment()
-            } else {
-                None
-            }
-        }) {
-            if std::mem::take(&mut first) {
-                write!(f, [line_suffix(&text("  "))])?;
-            }
-            write!(f, [line_suffix(&literal(range))])?;
-        }
+        write!(f, [end_of_line_comments(expr)])?;
 
         let magic_trailing_comma = expr.trivia.iter().any(|c| c.kind.is_magic_trailing_comma());
         write!(
@@ -394,14 +366,7 @@ fn format_call(
                     write!(
                         f,
                         [group(&format_args![&format_with(|f| {
-                            if let Some(arg) = &keyword.node.arg {
-                                write!(f, [dynamic_text(arg, TextSize::default())])?;
-                                write!(f, [text("=")])?;
-                                write!(f, [keyword.node.value.format()])?;
-                            } else {
-                                write!(f, [text("**")])?;
-                                write!(f, [keyword.node.value.format()])?;
-                            }
+                            write!(f, [keyword.format()])?;
                             Ok(())
                         })])]
                     )?;
@@ -736,19 +701,15 @@ fn format_named_expr(
 fn format_bool_op(
     f: &mut Formatter<ASTFormatContext<'_>>,
     expr: &Expr,
-    op: &Boolop,
+    ops: &[BoolOp],
     values: &[Expr],
 ) -> FormatResult<()> {
-    let mut first = true;
-    for value in values {
-        if std::mem::take(&mut first) {
-            write!(f, [group(&format_args![value.format()])])?;
-        } else {
-            write!(f, [soft_line_break_or_space()])?;
-            write!(f, [op.format()])?;
-            write!(f, [space()])?;
-            write!(f, [group(&format_args![value.format()])])?;
-        }
+    write!(f, [group(&format_args![values[0].format()])])?;
+    for (op, value) in ops.iter().zip(&values[1..]) {
+        write!(f, [soft_line_break_or_space()])?;
+        write!(f, [op.format()])?;
+        write!(f, [space()])?;
+        write!(f, [group(&format_args![value.format()])])?;
     }
     write!(f, [end_of_line_comments(expr)])?;
     Ok(())
@@ -851,7 +812,7 @@ impl Format<ASTFormatContext<'_>> for FormatExpr<'_> {
         write!(f, [leading_comments(self.item)])?;
 
         match &self.item.node {
-            ExprKind::BoolOp { op, values } => format_bool_op(f, self.item, op, values),
+            ExprKind::BoolOp { ops, values } => format_bool_op(f, self.item, ops, values),
             ExprKind::NamedExpr { target, value } => format_named_expr(f, self.item, target, value),
             ExprKind::BinOp { left, op, right } => format_bin_op(f, self.item, left, op, right),
             ExprKind::UnaryOp { op, operand } => format_unary_op(f, self.item, op, operand),