Make BoolOp its own located token (#3265)

Cargo.lock

@@ -2141,6 +2141,7 @@ dependencies = [
  "clap 4.1.6",
  "insta",
  "is-macro",
+ "itertools",
  "once_cell",
  "ruff_formatter",
  "ruff_python",

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) {
+    fn visit_bool_op(&mut self, bool_op: &'a mut BoolOp) {
-        walk_boolop(self, 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 } => {
+        ExprKind::BoolOp { ops, values } => {
-            visitor.visit_boolop(op);
+            for op in ops {
-            for expr in values {
+                visitor.visit_bool_op(op);
-                visitor.visit_expr(expr);
+            }
+            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,13 +59,13 @@ 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 {
+impl From<&rustpython_parser::ast::Boolop> for BoolOpKind {
-    fn from(op: rustpython_parser::ast::Boolop) -> Self {
+    fn from(op: &rustpython_parser::ast::Boolop) -> Self {
         match op {
             rustpython_parser::ast::Boolop::And => Self::And,
             rustpython_parser::ast::Boolop::Or => Self::Or,
@@ -71,6 +73,8 @@ impl From<rustpython_parser::ast::Boolop> for Boolop {
     }
 }
 
+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/boolop.rs

@@ -1,32 +0,0 @@
-use ruff_formatter::prelude::*;
-use ruff_formatter::write;
-
-use crate::context::ASTFormatContext;
-use crate::cst::Boolop;
-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,
-            [text(match boolop {
-                Boolop::And => "and",
-                Boolop::Or => "or",
-            })]
-        )?;
-        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(")")])?;
-
+        write!(f, [end_of_line_comments(expr)])?;
-        // 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))])?;
-        }
     } else {
         write!(f, [text("(")])?;
-
+        write!(f, [end_of_line_comments(expr)])?;
-        // 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))])?;
-        }
 
         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, [keyword.format()])?;
-                                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()])?;
-                            }
                             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;
+    write!(f, [group(&format_args![values[0].format()])])?;
-    for value in values {
+    for (op, value) in ops.iter().zip(&values[1..]) {
-        if std::mem::take(&mut first) {
+        write!(f, [soft_line_break_or_space()])?;
-            write!(f, [group(&format_args![value.format()])])?;
+        write!(f, [op.format()])?;
-        } else {
+        write!(f, [space()])?;
-            write!(f, [soft_line_break_or_space()])?;
+        write!(f, [group(&format_args![value.format()])])?;
-            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),

crates/ruff_python_formatter/src/format/keyword.rs

@@ -0,0 +1,40 @@
+use ruff_formatter::prelude::*;
+use ruff_formatter::write;
+use ruff_text_size::TextSize;
+
+use crate::context::ASTFormatContext;
+use crate::cst::Keyword;
+use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
+use crate::shared_traits::AsFormat;
+
+pub struct FormatKeyword<'a> {
+    item: &'a Keyword,
+}
+
+impl AsFormat<ASTFormatContext<'_>> for Keyword {
+    type Format<'a> = FormatKeyword<'a>;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatKeyword { item: self }
+    }
+}
+
+impl Format<ASTFormatContext<'_>> for FormatKeyword<'_> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext>) -> FormatResult<()> {
+        let keyword = self.item;
+
+        write!(f, [leading_comments(keyword)])?;
+        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, [end_of_line_comments(keyword)])?;
+        write!(f, [trailing_comments(keyword)])?;
+
+        Ok(())
+    }
+}

crates/ruff_python_formatter/src/format/mod.rs

@@ -1,7 +1,7 @@
 mod alias;
 mod arg;
 mod arguments;
-mod boolop;
+mod bool_op;
 pub mod builders;
 mod cmpop;
 mod comments;
@@ -9,6 +9,7 @@ mod comprehension;
 mod excepthandler;
 mod expr;
 mod helpers;
+mod keyword;
 mod match_case;
 mod numbers;
 mod operator;

crates/ruff_python_formatter/src/format/stmt.rs

@@ -131,18 +131,7 @@ fn format_class_def(
                 }
 
                 for (i, keyword) in keywords.iter().enumerate() {
-                    if let Some(arg) = &keyword.node.arg {
+                    write!(f, [keyword.format()])?;
-                        write!(
-                            f,
-                            [
-                                dynamic_text(arg, TextSize::default()),
-                                text("="),
-                                keyword.node.value.format()
-                            ]
-                        )?;
-                    } else {
-                        write!(f, [text("**"), keyword.node.value.format()])?;
-                    }
                     if i < keywords.len() - 1 {
                         write!(f, [text(","), soft_line_break_or_space()])?;
                     } else {

crates/ruff_python_formatter/src/newlines.rs

@@ -2,7 +2,10 @@ use rustpython_parser::ast::Constant;
 
 use crate::core::visitor;
 use crate::core::visitor::Visitor;
-use crate::cst::{ExcepthandlerKind, Expr, ExprKind, Stmt, StmtKind};
+use crate::cst::{
+    Alias, Arg, BoolOp, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Keyword, Pattern,
+    SliceIndex, Stmt, StmtKind,
+};
 use crate::trivia::{Relationship, Trivia, TriviaKind};
 
 #[derive(Debug, Copy, Clone)]
@@ -296,9 +299,43 @@ struct ExprNormalizer;
 impl<'a> Visitor<'a> for ExprNormalizer {
     fn visit_expr(&mut self, expr: &'a mut Expr) {
         expr.trivia.retain(|c| !c.kind.is_empty_line());
-
         visitor::walk_expr(self, expr);
     }
+
+    fn visit_alias(&mut self, alias: &'a mut Alias) {
+        alias.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_alias(self, alias);
+    }
+
+    fn visit_arg(&mut self, arg: &'a mut Arg) {
+        arg.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_arg(self, arg);
+    }
+
+    fn visit_excepthandler(&mut self, excepthandler: &'a mut Excepthandler) {
+        excepthandler.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_excepthandler(self, excepthandler);
+    }
+
+    fn visit_keyword(&mut self, keyword: &'a mut Keyword) {
+        keyword.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_keyword(self, keyword);
+    }
+
+    fn visit_bool_op(&mut self, bool_op: &'a mut BoolOp) {
+        bool_op.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_bool_op(self, bool_op);
+    }
+
+    fn visit_slice_index(&mut self, slice_index: &'a mut SliceIndex) {
+        slice_index.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_slice_index(self, slice_index);
+    }
+
+    fn visit_pattern(&mut self, pattern: &'a mut Pattern) {
+        pattern.trivia.retain(|c| !c.kind.is_empty_line());
+        visitor::walk_pattern(self, pattern);
+    }
 }
 
 pub fn normalize_newlines(python_cst: &mut [Stmt]) {

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap

@@ -178,7 +178,7 @@ instruction()#comment with bad spacing
 ```diff
 --- Black
 +++ Ruff
-@@ -72,14 +72,20 @@
+@@ -72,7 +72,11 @@
              body,
              parameters.children[-1],  # )2
          ]
@@ -190,21 +190,8 @@ instruction()#comment with bad spacing
 +        ]  # type: ignore
      if (
          self._proc is not None
--        # has the child process finished?
+         # has the child process finished?
--        and self._returncode is None
+@@ -103,35 +107,35 @@
--        # the child process has finished, but the
-+        and # has the child process finished?
-+        self._returncode
-+        is None
-+        and # the child process has finished, but the
-         # transport hasn't been notified yet?
--        and self._proc.poll() is None
-+        self._proc.poll()
-+        is None
-     ):
-         pass
-     # no newline before or after
-@@ -103,35 +109,35 @@
      ############################################################################
  
      call2(
@@ -219,10 +206,8 @@ instruction()#comment with bad spacing
          """
  short
  """,
--        # yup
+         # yup
--        arg3=True,
+         arg3=True,
-+        arg3=# yup
-+        True,
      )
 -    lcomp = [
 -        element for element in collection if element is not None  # yup  # yup  # right
@@ -256,7 +241,7 @@ instruction()#comment with bad spacing
      ]
      while True:
          if False:
-@@ -167,7 +173,7 @@
+@@ -167,7 +171,7 @@
  #######################
  
  
@@ -351,13 +336,11 @@ def inline_comments_in_brackets_ruin_everything():
         ]  # type: ignore
     if (
         self._proc is not None
-        and # has the child process finished?
+        # has the child process finished?
-        self._returncode
+        and self._returncode is None
-        is None
+        # the child process has finished, but the
-        and # the child process has finished, but the
         # transport hasn't been notified yet?
-        self._proc.poll()
+        and self._proc.poll() is None
-        is None
     ):
         pass
     # no newline before or after
@@ -389,8 +372,8 @@ short
         """
 short
 """,
-        arg3=# yup
+        # yup
-        True,
+        arg3=True,
     )
     lcomp = [element for element in collection if element is not None]  # yup  # yup  # right
     lcomp2 = [

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments3_py.snap

@@ -77,30 +77,6 @@ def func():
      ]
      # Capture each of the exceptions in the MultiError along with each of their causes and contexts
      if isinstance(exc_value, MultiError):
-@@ -26,9 +27,9 @@
-                         limit=limit,
-                         lookup_lines=lookup_lines,
-                         capture_locals=capture_locals,
--                        # copy the set of _seen exceptions so that duplicates
-+                        _seen=# copy the set of _seen exceptions so that duplicates
-                         # shared between sub-exceptions are not omitted
--                        _seen=set(_seen),
-+                        set(_seen),
-                     )
-                     # This should be left alone (after)
-                 )
-@@ -39,9 +40,9 @@
-         limit=limit,
-         lookup_lines=lookup_lines,
-         capture_locals=capture_locals,
--        # copy the set of _seen exceptions so that duplicates
-+        _seen=# copy the set of _seen exceptions so that duplicates
-         # shared between sub-exceptions are not omitted
--        _seen=set(_seen),
-+        set(_seen),
-     )
- 
- 
 ```
 
 ## Ruff Output
@@ -135,9 +111,9 @@ def func():
                         limit=limit,
                         lookup_lines=lookup_lines,
                         capture_locals=capture_locals,
-                        _seen=# copy the set of _seen exceptions so that duplicates
+                        # copy the set of _seen exceptions so that duplicates
                         # shared between sub-exceptions are not omitted
-                        set(_seen),
+                        _seen=set(_seen),
                     )
                     # This should be left alone (after)
                 )
@@ -148,9 +124,9 @@ def func():
         limit=limit,
         lookup_lines=lookup_lines,
         capture_locals=capture_locals,
-        _seen=# copy the set of _seen exceptions so that duplicates
+        # copy the set of _seen exceptions so that duplicates
         # shared between sub-exceptions are not omitted
-        set(_seen),
+        _seen=set(_seen),
     )
 
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__composition_no_trailing_comma_py.snap

@@ -204,15 +204,9 @@ class C:
              )
              self.assertEqual(
                  unstyle(str(report)),
-@@ -22,133 +23,155 @@
+@@ -25,11 +26,8 @@
-                 if (
+                     # Rule 2
-                     # Rule 1
+                     and i % 3 == 0
-                     i % 2 == 0
--                    # Rule 2
--                    and i % 3 == 0
-+                    and # Rule 2
-+                    i % 3
-+                    == 0
                  ):
 -                    while (
 -                        # Just a comment
@@ -224,13 +218,9 @@ class C:
                          print(i)
          xxxxxxxxxxxxxxxx = Yyyy2YyyyyYyyyyy(
              push_manager=context.request.resource_manager,
-             max_items_to_push=num_items,
+@@ -39,116 +37,140 @@
-             batch_size=Yyyy2YyyyYyyyyYyyy.FULL_SIZE,
+             # Only send the first n items.
-         ).push(
+             items=items[:num_items]
--            # Only send the first n items.
--            items=items[:num_items]
-+            items=# Only send the first n items.
-+            items[:num_items]
          )
 -        return (
 -            'Utterly failed doctest test for %s\n   File "%s", line %s, in %s\n\n%s'
@@ -460,7 +450,7 @@ class C:
              "Not what we expected and the message is too long to fit in one line"
              " because it's too long"
          )
-@@ -161,9 +184,8 @@
+@@ -161,9 +183,8 @@
                        8 STORE_ATTR               0 (x)
                       10 LOAD_CONST               0 (None)
                       12 RETURN_VALUE
@@ -502,9 +492,8 @@ class C:
                 if (
                     # Rule 1
                     i % 2 == 0
-                    and # Rule 2
+                    # Rule 2
-                    i % 3
+                    and i % 3 == 0
-                    == 0
                 ):
                     while # Just a comment
                     call():
@@ -514,8 +503,8 @@ class C:
             max_items_to_push=num_items,
             batch_size=Yyyy2YyyyYyyyyYyyy.FULL_SIZE,
         ).push(
-            items=# Only send the first n items.
+            # Only send the first n items.
-            items[:num_items]
+            items=items[:num_items]
         )
         return 'Utterly failed doctest test for %s\n   File "%s", line %s, in %s\n\n%s'
         % (test.name, test.filename, lineno, lname, err)

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__composition_py.snap

@@ -204,15 +204,9 @@ class C:
              )
              self.assertEqual(
                  unstyle(str(report)),
-@@ -22,133 +23,155 @@
+@@ -25,11 +26,8 @@
-                 if (
+                     # Rule 2
-                     # Rule 1
+                     and i % 3 == 0
-                     i % 2 == 0
--                    # Rule 2
--                    and i % 3 == 0
-+                    and # Rule 2
-+                    i % 3
-+                    == 0
                  ):
 -                    while (
 -                        # Just a comment
@@ -224,13 +218,9 @@ class C:
                          print(i)
          xxxxxxxxxxxxxxxx = Yyyy2YyyyyYyyyyy(
              push_manager=context.request.resource_manager,
-             max_items_to_push=num_items,
+@@ -39,116 +37,140 @@
-             batch_size=Yyyy2YyyyYyyyyYyyy.FULL_SIZE,
+             # Only send the first n items.
-         ).push(
+             items=items[:num_items]
--            # Only send the first n items.
--            items=items[:num_items]
-+            items=# Only send the first n items.
-+            items[:num_items]
          )
 -        return (
 -            'Utterly failed doctest test for %s\n   File "%s", line %s, in %s\n\n%s'
@@ -460,7 +450,7 @@ class C:
              "Not what we expected and the message is too long to fit in one line"
              " because it's too long"
          )
-@@ -161,9 +184,8 @@
+@@ -161,9 +183,8 @@
                        8 STORE_ATTR               0 (x)
                       10 LOAD_CONST               0 (None)
                       12 RETURN_VALUE
@@ -502,9 +492,8 @@ class C:
                 if (
                     # Rule 1
                     i % 2 == 0
-                    and # Rule 2
+                    # Rule 2
-                    i % 3
+                    and i % 3 == 0
-                    == 0
                 ):
                     while # Just a comment
                     call():
@@ -514,8 +503,8 @@ class C:
             max_items_to_push=num_items,
             batch_size=Yyyy2YyyyYyyyyYyyy.FULL_SIZE,
         ).push(
-            items=# Only send the first n items.
+            # Only send the first n items.
-            items[:num_items]
+            items=items[:num_items]
         )
         return 'Utterly failed doctest test for %s\n   File "%s", line %s, in %s\n\n%s'
         % (test.name, test.filename, lineno, lname, err)

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtonoff_py.snap

@@ -394,25 +394,8 @@ d={'a':1,
          # fmt: on
      )
  
-@@ -200,8 +213,8 @@
+@@ -217,8 +230,7 @@
-     xxxxxx_xxxxxx=2,
+     xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=5,
-     xxxxxx_xxxxx_xxxxxxxx=70,
-     xxxxxx_xxxxxx_xxxxx=True,
--    # fmt: off
--    xxxxxxx_xxxxxxxxxxxx={
-+    xxxxxxx_xxxxxxxxxxxx=# fmt: off
-+    {
-         "xxxxxxxx": {
-             "xxxxxx": False,
-             "xxxxxxx": False,
-@@ -213,12 +226,11 @@
-             "xxxx_xxxxxx": "xxxxxx",
-         },
-     },
--    # fmt: on
--    xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=5,
-+    xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=# fmt: on
-+    5,
  )
  # fmt: off
 -yield  'hello'
@@ -643,8 +626,8 @@ cfg.rule(
     xxxxxx_xxxxxx=2,
     xxxxxx_xxxxx_xxxxxxxx=70,
     xxxxxx_xxxxxx_xxxxx=True,
-    xxxxxxx_xxxxxxxxxxxx=# fmt: off
+    # fmt: off
-    {
+    xxxxxxx_xxxxxxxxxxxx={
         "xxxxxxxx": {
             "xxxxxx": False,
             "xxxxxxx": False,
@@ -656,8 +639,8 @@ cfg.rule(
             "xxxx_xxxxxx": "xxxxxx",
         },
     },
-    xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=# fmt: on
+    # fmt: on
-    5,
+    xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=5,
 )
 # fmt: off
 yield "hello"

crates/ruff_python_formatter/src/trivia.rs

@@ -5,35 +5,39 @@ use rustpython_parser::Tok;
 
 use crate::core::types::Range;
 use crate::cst::{
-    Alias, Arg, Body, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Pattern, PatternKind,
+    Alias, Arg, Body, BoolOp, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Keyword, Pattern,
-    SliceIndex, SliceIndexKind, Stmt, StmtKind,
+    PatternKind, SliceIndex, SliceIndexKind, Stmt, StmtKind,
 };
 
 #[derive(Clone, Debug)]
 pub enum Node<'a> {
-    Mod(&'a [Stmt]),
-    Body(&'a Body),
-    Stmt(&'a Stmt),
-    Expr(&'a Expr),
     Alias(&'a Alias),
     Arg(&'a Arg),
+    Body(&'a Body),
+    BoolOp(&'a BoolOp),
     Excepthandler(&'a Excepthandler),
-    SliceIndex(&'a SliceIndex),
+    Expr(&'a Expr),
+    Keyword(&'a Keyword),
+    Mod(&'a [Stmt]),
     Pattern(&'a Pattern),
+    SliceIndex(&'a SliceIndex),
+    Stmt(&'a Stmt),
 }
 
 impl Node<'_> {
     pub fn id(&self) -> usize {
         match self {
-            Node::Mod(nodes) => nodes as *const _ as usize,
-            Node::Body(node) => node.id(),
-            Node::Stmt(node) => node.id(),
-            Node::Expr(node) => node.id(),
             Node::Alias(node) => node.id(),
             Node::Arg(node) => node.id(),
+            Node::Body(node) => node.id(),
+            Node::BoolOp(node) => node.id(),
             Node::Excepthandler(node) => node.id(),
-            Node::SliceIndex(node) => node.id(),
+            Node::Expr(node) => node.id(),
+            Node::Keyword(node) => node.id(),
+            Node::Mod(nodes) => nodes as *const _ as usize,
             Node::Pattern(node) => node.id(),
+            Node::SliceIndex(node) => node.id(),
+            Node::Stmt(node) => node.id(),
         }
     }
 }
@@ -236,6 +240,7 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                 result.push(Node::Stmt(stmt));
             }
         }
+        Node::BoolOp(..) => {}
         Node::Stmt(stmt) => match &stmt.node {
             StmtKind::Return { value } => {
                 if let Some(value) = value {
@@ -309,7 +314,7 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                     result.push(Node::Expr(base));
                 }
                 for keyword in keywords {
-                    result.push(Node::Expr(&keyword.node.value));
+                    result.push(Node::Keyword(keyword));
                 }
                 result.push(Node::Body(body));
             }
@@ -448,8 +453,10 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
             }
         }
         Node::Expr(expr) => match &expr.node {
-            ExprKind::BoolOp { values, .. } => {
+            ExprKind::BoolOp { ops, values } => {
-                for value in values {
+                result.push(Node::Expr(&values[0]));
+                for (op, value) in ops.iter().zip(&values[1..]) {
+                    result.push(Node::BoolOp(op));
                     result.push(Node::Expr(value));
                 }
             }
@@ -565,7 +572,7 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                     result.push(Node::Expr(arg));
                 }
                 for keyword in keywords {
-                    result.push(Node::Expr(&keyword.node.value));
+                    result.push(Node::Keyword(keyword));
                 }
             }
             ExprKind::FormattedValue {
@@ -612,6 +619,9 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                 }
             }
         },
+        Node::Keyword(keyword) => {
+            result.push(Node::Expr(&keyword.node.value));
+        }
         Node::Alias(..) => {}
         Node::Excepthandler(excepthandler) => {
             let ExcepthandlerKind::ExceptHandler { type_, body, .. } = &excepthandler.node;
@@ -703,52 +713,60 @@ pub fn decorate_token<'a>(
         let middle = (left + right) / 2;
         let child = &child_nodes[middle];
         let start = match &child {
-            Node::Body(node) => node.location,
-            Node::Stmt(node) => node.location,
-            Node::Expr(node) => node.location,
             Node::Alias(node) => node.location,
             Node::Arg(node) => node.location,
+            Node::Body(node) => node.location,
+            Node::BoolOp(node) => node.location,
             Node::Excepthandler(node) => node.location,
-            Node::SliceIndex(node) => node.location,
+            Node::Expr(node) => node.location,
-            Node::Pattern(node) => node.location,
+            Node::Keyword(node) => node.location,
             Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
+            Node::Pattern(node) => node.location,
+            Node::SliceIndex(node) => node.location,
+            Node::Stmt(node) => node.location,
         };
         let end = match &child {
-            Node::Body(node) => node.end_location.unwrap(),
-            Node::Stmt(node) => node.end_location.unwrap(),
-            Node::Expr(node) => node.end_location.unwrap(),
             Node::Alias(node) => node.end_location.unwrap(),
             Node::Arg(node) => node.end_location.unwrap(),
+            Node::Body(node) => node.end_location.unwrap(),
+            Node::BoolOp(node) => node.end_location.unwrap(),
             Node::Excepthandler(node) => node.end_location.unwrap(),
-            Node::SliceIndex(node) => node.end_location.unwrap(),
+            Node::Expr(node) => node.end_location.unwrap(),
-            Node::Pattern(node) => node.end_location.unwrap(),
+            Node::Keyword(node) => node.end_location.unwrap(),
             Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
+            Node::Pattern(node) => node.end_location.unwrap(),
+            Node::SliceIndex(node) => node.end_location.unwrap(),
+            Node::Stmt(node) => node.end_location.unwrap(),
         };
 
         if let Some(existing) = &enclosed_node {
             // Special-case: if we're dealing with a statement that's a single expression,
             // we want to treat the expression as the enclosed node.
             let existing_start = match &existing {
-                Node::Body(node) => node.location,
-                Node::Stmt(node) => node.location,
-                Node::Expr(node) => node.location,
                 Node::Alias(node) => node.location,
                 Node::Arg(node) => node.location,
+                Node::Body(node) => node.location,
+                Node::BoolOp(node) => node.location,
                 Node::Excepthandler(node) => node.location,
-                Node::SliceIndex(node) => node.location,
+                Node::Expr(node) => node.location,
-                Node::Pattern(node) => node.location,
+                Node::Keyword(node) => node.location,
                 Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
+                Node::Pattern(node) => node.location,
+                Node::SliceIndex(node) => node.location,
+                Node::Stmt(node) => node.location,
             };
             let existing_end = match &existing {
-                Node::Body(node) => node.end_location.unwrap(),
-                Node::Stmt(node) => node.end_location.unwrap(),
-                Node::Expr(node) => node.end_location.unwrap(),
                 Node::Alias(node) => node.end_location.unwrap(),
                 Node::Arg(node) => node.end_location.unwrap(),
+                Node::Body(node) => node.end_location.unwrap(),
+                Node::BoolOp(node) => node.end_location.unwrap(),
                 Node::Excepthandler(node) => node.end_location.unwrap(),
-                Node::SliceIndex(node) => node.end_location.unwrap(),
+                Node::Expr(node) => node.end_location.unwrap(),
-                Node::Pattern(node) => node.end_location.unwrap(),
+                Node::Keyword(node) => node.end_location.unwrap(),
                 Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
+                Node::Pattern(node) => node.end_location.unwrap(),
+                Node::SliceIndex(node) => node.end_location.unwrap(),
+                Node::Stmt(node) => node.end_location.unwrap(),
             };
             if start == existing_start && end == existing_end {
                 enclosed_node = Some(child.clone());
@@ -803,40 +821,20 @@ pub fn decorate_token<'a>(
 
 #[derive(Debug, Default)]
 pub struct TriviaIndex {
-    pub body: FxHashMap<usize, Vec<Trivia>>,
-    pub stmt: FxHashMap<usize, Vec<Trivia>>,
-    pub expr: FxHashMap<usize, Vec<Trivia>>,
     pub alias: FxHashMap<usize, Vec<Trivia>>,
     pub arg: FxHashMap<usize, Vec<Trivia>>,
+    pub body: FxHashMap<usize, Vec<Trivia>>,
+    pub bool_op: FxHashMap<usize, Vec<Trivia>>,
     pub excepthandler: FxHashMap<usize, Vec<Trivia>>,
-    pub slice_index: FxHashMap<usize, Vec<Trivia>>,
+    pub expr: FxHashMap<usize, Vec<Trivia>>,
+    pub keyword: FxHashMap<usize, Vec<Trivia>>,
     pub pattern: FxHashMap<usize, Vec<Trivia>>,
+    pub slice_index: FxHashMap<usize, Vec<Trivia>>,
+    pub stmt: FxHashMap<usize, Vec<Trivia>>,
 }
 
 fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
     match node {
-        Node::Mod(_) => {}
-        Node::Body(node) => {
-            trivia
-                .body
-                .entry(node.id())
-                .or_insert_with(Vec::new)
-                .push(comment);
-        }
-        Node::Stmt(node) => {
-            trivia
-                .stmt
-                .entry(node.id())
-                .or_insert_with(Vec::new)
-                .push(comment);
-        }
-        Node::Expr(node) => {
-            trivia
-                .expr
-                .entry(node.id())
-                .or_insert_with(Vec::new)
-                .push(comment);
-        }
         Node::Alias(node) => {
             trivia
                 .alias
@@ -851,6 +849,20 @@ fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
                 .or_insert_with(Vec::new)
                 .push(comment);
         }
+        Node::Body(node) => {
+            trivia
+                .body
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
+        Node::BoolOp(node) => {
+            trivia
+                .bool_op
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
         Node::Excepthandler(node) => {
             trivia
                 .excepthandler
@@ -858,9 +870,16 @@ fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
                 .or_insert_with(Vec::new)
                 .push(comment);
         }
-        Node::SliceIndex(node) => {
+        Node::Expr(node) => {
             trivia
-                .slice_index
+                .expr
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
+        Node::Keyword(node) => {
+            trivia
+                .keyword
                 .entry(node.id())
                 .or_insert_with(Vec::new)
                 .push(comment);
@@ -872,6 +891,21 @@ fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
                 .or_insert_with(Vec::new)
                 .push(comment);
         }
+        Node::SliceIndex(node) => {
+            trivia
+                .slice_index
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
+        Node::Stmt(node) => {
+            trivia
+                .stmt
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
+        Node::Mod(_) => {}
     }
 }