Add builders for common comment rendering (#3232)

crates/ruff_python_formatter/src/builders.rs

@@ -1,77 +0,0 @@
-use ruff_formatter::prelude::*;
-use ruff_formatter::{write, Format};
-use ruff_text_size::TextRange;
-
-use crate::context::ASTFormatContext;
-use crate::core::types::Range;
-use crate::trivia::{Trivia, TriviaKind};
-
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
-pub struct Literal {
-    range: Range,
-}
-
-impl Format<ASTFormatContext<'_>> for Literal {
-    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
-        let (text, start, end) = f.context().locator().slice(self.range);
-        f.write_element(FormatElement::StaticTextSlice {
-            text,
-            range: TextRange::new(start.try_into().unwrap(), end.try_into().unwrap()),
-        })
-    }
-}
-
-// TODO(charlie): We still can't use this everywhere we'd like. We need the AST
-// to include ranges for all `Ident` nodes.
-#[inline]
-pub const fn literal(range: Range) -> Literal {
-    Literal { range }
-}
-
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
-pub struct LeadingComments<'a> {
-    comments: &'a [Trivia],
-}
-
-impl Format<ASTFormatContext<'_>> for LeadingComments<'_> {
-    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
-        for comment in self.comments {
-            if comment.relationship.is_leading() {
-                if let TriviaKind::OwnLineComment(range) = comment.kind {
-                    write!(f, [hard_line_break()])?;
-                    write!(f, [literal(range)])?;
-                }
-            }
-        }
-        Ok(())
-    }
-}
-
-#[inline]
-pub const fn leading_comments(comments: &[Trivia]) -> LeadingComments {
-    LeadingComments { comments }
-}
-
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
-pub struct TrailingComments<'a> {
-    comments: &'a [Trivia],
-}
-
-impl Format<ASTFormatContext<'_>> for TrailingComments<'_> {
-    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
-        for comment in self.comments {
-            if comment.relationship.is_trailing() {
-                if let TriviaKind::OwnLineComment(range) = comment.kind {
-                    write!(f, [hard_line_break()])?;
-                    write!(f, [literal(range)])?;
-                }
-            }
-        }
-        Ok(())
-    }
-}
-
-#[inline]
-pub const fn trailing_comments(comments: &[Trivia]) -> TrailingComments {
-    TrailingComments { comments }
-}

crates/ruff_python_formatter/src/format/alias.rs

@@ -2,9 +2,9 @@ use ruff_formatter::prelude::*;
 use ruff_formatter::write;
 use ruff_text_size::TextSize;
 
-use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::cst::Alias;
+use crate::format::comments::end_of_line_comments;
 use crate::shared_traits::AsFormat;
 
 pub struct FormatAlias<'a> {
@@ -29,20 +29,7 @@ impl Format<ASTFormatContext<'_>> for FormatAlias<'_> {
             write!(f, [dynamic_text(asname, TextSize::default())])?;
         }
 
-        // Format any end-of-line comments.
+        write!(f, [end_of_line_comments(alias)])?;
-        let mut first = true;
-        for range in alias.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))])?;
-        }
 
         Ok(())
     }

crates/ruff_python_formatter/src/format/builders.rs

@@ -1,8 +1,9 @@
 use ruff_formatter::prelude::*;
 use ruff_formatter::{write, Format};
-use ruff_text_size::TextSize;
+use ruff_text_size::{TextRange, TextSize};
 
 use crate::context::ASTFormatContext;
+use crate::core::types::Range;
 use crate::cst::Stmt;
 use crate::shared_traits::AsFormat;
 
@@ -28,6 +29,26 @@ pub fn block(body: &[Stmt]) -> Block {
     Block { body }
 }
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct Literal {
+    range: Range,
+}
+
+impl Format<ASTFormatContext<'_>> for Literal {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        let (text, start, end) = f.context().locator().slice(self.range);
+        f.write_element(FormatElement::StaticTextSlice {
+            text,
+            range: TextRange::new(start.try_into().unwrap(), end.try_into().unwrap()),
+        })
+    }
+}
+
+#[inline]
+pub const fn literal(range: Range) -> Literal {
+    Literal { range }
+}
+
 pub(crate) const fn join_names(names: &[String]) -> JoinNames {
     JoinNames { names }
 }

crates/ruff_python_formatter/src/format/comments.rs

@@ -0,0 +1,119 @@
+use ruff_formatter::prelude::*;
+use ruff_formatter::{write, Format};
+
+use crate::context::ASTFormatContext;
+use crate::cst::Located;
+use crate::format::builders::literal;
+use crate::trivia::TriviaKind;
+
+#[derive(Debug)]
+pub struct LeadingComments<'a, T> {
+    item: &'a Located<T>,
+}
+
+impl<T> Format<ASTFormatContext<'_>> for LeadingComments<'_, T> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        for trivia in &self.item.trivia {
+            if trivia.relationship.is_leading() {
+                match trivia.kind {
+                    TriviaKind::EmptyLine => {
+                        write!(f, [empty_line()])?;
+                    }
+                    TriviaKind::OwnLineComment(range) => {
+                        write!(f, [literal(range), hard_line_break()])?;
+                    }
+                    _ => {}
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn leading_comments<T>(item: &Located<T>) -> LeadingComments<'_, T> {
+    LeadingComments { item }
+}
+
+#[derive(Debug)]
+pub struct TrailingComments<'a, T> {
+    item: &'a Located<T>,
+}
+
+impl<T> Format<ASTFormatContext<'_>> for TrailingComments<'_, T> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        for trivia in &self.item.trivia {
+            if trivia.relationship.is_trailing() {
+                match trivia.kind {
+                    TriviaKind::EmptyLine => {
+                        write!(f, [empty_line()])?;
+                    }
+                    TriviaKind::OwnLineComment(range) => {
+                        write!(f, [literal(range), hard_line_break()])?;
+                    }
+                    _ => {}
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn trailing_comments<T>(item: &Located<T>) -> TrailingComments<'_, T> {
+    TrailingComments { item }
+}
+
+#[derive(Debug)]
+pub struct EndOfLineComments<'a, T> {
+    item: &'a Located<T>,
+}
+
+impl<T> Format<ASTFormatContext<'_>> for EndOfLineComments<'_, T> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        let mut first = true;
+        for range in self.item.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))])?;
+        }
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn end_of_line_comments<T>(item: &Located<T>) -> EndOfLineComments<'_, T> {
+    EndOfLineComments { item }
+}
+
+#[derive(Debug)]
+pub struct DanglingComments<'a, T> {
+    item: &'a Located<T>,
+}
+
+impl<T> Format<ASTFormatContext<'_>> for DanglingComments<'_, T> {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        for trivia in &self.item.trivia {
+            if trivia.relationship.is_dangling() {
+                if let TriviaKind::OwnLineComment(range) = trivia.kind {
+                    write!(f, [hard_line_break()])?;
+                    write!(f, [literal(range)])?;
+                    write!(f, [hard_line_break()])?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn dangling_comments<T>(item: &Located<T>) -> DanglingComments<'_, T> {
+    DanglingComments { item }
+}

crates/ruff_python_formatter/src/format/excepthandler.rs

@@ -2,10 +2,10 @@ use ruff_formatter::prelude::*;
 use ruff_formatter::write;
 use ruff_text_size::TextSize;
 
-use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::cst::{Excepthandler, ExcepthandlerKind};
 use crate::format::builders::block;
+use crate::format::comments::end_of_line_comments;
 use crate::shared_traits::AsFormat;
 
 pub struct FormatExcepthandler<'a> {
@@ -41,21 +41,7 @@ impl Format<ASTFormatContext<'_>> for FormatExcepthandler<'_> {
             }
         }
         write!(f, [text(":")])?;
-
+        write!(f, [end_of_line_comments(excepthandler)])?;
-        // Format any end-of-line comments.
-        let mut first = true;
-        for range in excepthandler.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, [block_indent(&block(body))])?;
 

crates/ruff_python_formatter/src/format/expr.rs

@@ -6,13 +6,14 @@ use ruff_formatter::prelude::*;
 use ruff_formatter::{format_args, write};
 use ruff_text_size::TextSize;
 
-use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::core::types::Range;
 use crate::cst::{
     Arguments, Boolop, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, SliceIndex,
     SliceIndexKind, Unaryop,
 };
+use crate::format::builders::literal;
+use crate::format::comments::{dangling_comments, end_of_line_comments, leading_comments};
 use crate::format::helpers::{is_self_closing, is_simple_power, is_simple_slice};
 use crate::format::numbers::{complex_literal, float_literal, int_literal};
 use crate::format::strings::string_literal;
@@ -29,22 +30,7 @@ fn format_starred(
     value: &Expr,
 ) -> FormatResult<()> {
     write!(f, [text("*"), value.format()])?;
-
+    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))])?;
-    }
-
     Ok(())
 }
 
@@ -54,22 +40,7 @@ fn format_name(
     _id: &str,
 ) -> FormatResult<()> {
     write!(f, [literal(Range::from_located(expr))])?;
-
+    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))])?;
-    }
-
     Ok(())
 }
 
@@ -222,17 +193,7 @@ fn format_slice(
                 write!(f, [value.format()])?;
             }
 
-            // Apply any dangling comments.
+            write!(f, [dangling_comments(lower)])?;
-            for trivia in &lower.trivia {
-                if trivia.relationship.is_dangling() {
-                    if let TriviaKind::OwnLineComment(range) = trivia.kind {
-                        write!(f, [expand_parent()])?;
-                        write!(f, [hard_line_break()])?;
-                        write!(f, [literal(range)])?;
-                        write!(f, [hard_line_break()])?;
-                    }
-                }
-            }
 
             if matches!(lower.node, SliceIndexKind::Index { .. }) {
                 if !is_simple {
@@ -240,21 +201,7 @@ fn format_slice(
                 }
             }
             write!(f, [text(":")])?;
-
+            write!(f, [end_of_line_comments(lower)])?;
-            // Format any end-of-line comments.
-            let mut first = true;
-            for range in lower.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))])?;
-            }
 
             if let SliceIndexKind::Index { value } = &upper.node {
                 if !is_simple {
@@ -264,32 +211,8 @@ fn format_slice(
                 write!(f, [value.format()])?;
             }
 
-            // Apply any dangling comments.
+            write!(f, [dangling_comments(upper)])?;
-            for trivia in &upper.trivia {
+            write!(f, [end_of_line_comments(upper)])?;
-                if trivia.relationship.is_dangling() {
-                    if let TriviaKind::OwnLineComment(range) = trivia.kind {
-                        write!(f, [expand_parent()])?;
-                        write!(f, [hard_line_break()])?;
-                        write!(f, [literal(range)])?;
-                        write!(f, [hard_line_break()])?;
-                    }
-                }
-            }
-
-            // Format any end-of-line comments.
-            let mut first = true;
-            for range in upper.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))])?;
-            }
 
             if let Some(step) = step {
                 if matches!(upper.node, SliceIndexKind::Index { .. }) {
@@ -307,51 +230,14 @@ fn format_slice(
                     write!(f, [value.format()])?;
                 }
 
-                // Apply any dangling comments.
+                write!(f, [dangling_comments(step)])?;
-                for trivia in &step.trivia {
+                write!(f, [end_of_line_comments(step)])?;
-                    if trivia.relationship.is_dangling() {
-                        if let TriviaKind::OwnLineComment(range) = trivia.kind {
-                            write!(f, [expand_parent()])?;
-                            write!(f, [hard_line_break()])?;
-                            write!(f, [literal(range)])?;
-                            write!(f, [hard_line_break()])?;
-                        }
-                    }
-                }
-
-                // Format any end-of-line comments.
-                let mut first = true;
-                for range in step.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))])?;
-                }
             }
             Ok(())
         }))]
     )?;
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(expr)])?;
-    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))])?;
-    }
 
     Ok(())
 }
@@ -513,16 +399,7 @@ fn format_call(
                     }
                 }
 
-                // Apply any dangling trailing comments.
+                write!(f, [dangling_comments(expr)])?;
-                for trivia in &expr.trivia {
-                    if trivia.relationship.is_dangling() {
-                        if let TriviaKind::OwnLineComment(range) = trivia.kind {
-                            write!(f, [expand_parent()])?;
-                            write!(f, [hard_line_break()])?;
-                            write!(f, [literal(range)])?;
-                        }
-                    }
-                }
 
                 Ok(())
             }))])]
@@ -709,20 +586,7 @@ fn format_compare(
         write!(f, [group(&format_args![comparators[i].format()])])?;
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(expr)])?;
-    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))])?;
-    }
 
     Ok(())
 }
@@ -834,20 +698,7 @@ fn format_attribute(
     write!(f, [text(".")])?;
     write!(f, [dynamic_text(attr, TextSize::default())])?;
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(expr)])?;
-    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))])?;
-    }
 
     Ok(())
 }
@@ -870,20 +721,7 @@ fn format_bool_op(
         }
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(expr)])?;
-    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))])?;
-    }
 
     Ok(())
 }
@@ -908,20 +746,7 @@ fn format_bin_op(
     }
     write!(f, [group(&format_args![right.format()])])?;
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(expr)])?;
-    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))])?;
-    }
 
     Ok(())
 }
@@ -996,16 +821,7 @@ impl Format<ASTFormatContext<'_>> for FormatExpr<'_> {
             write!(f, [text("(")])?;
         }
 
-        // Any leading comments come on the line before.
+        write!(f, [leading_comments(self.item)])?;
-        for trivia in &self.item.trivia {
-            if trivia.relationship.is_leading() {
-                if let TriviaKind::OwnLineComment(range) = trivia.kind {
-                    write!(f, [expand_parent()])?;
-                    write!(f, [literal(range)])?;
-                    write!(f, [hard_line_break()])?;
-                }
-            }
-        }
 
         match &self.item.node {
             ExprKind::BoolOp { op, values } => format_bool_op(f, self.item, op, values),

crates/ruff_python_formatter/src/format/mod.rs

@@ -4,6 +4,7 @@ mod arguments;
 mod boolop;
 pub mod builders;
 mod cmpop;
+mod comments;
 mod comprehension;
 mod excepthandler;
 mod expr;

crates/ruff_python_formatter/src/format/numbers.rs

@@ -1,11 +1,12 @@
+use rustpython_parser::ast::Location;
+
 use ruff_formatter::prelude::*;
 use ruff_formatter::{write, Format};
 use ruff_text_size::TextSize;
-use rustpython_parser::ast::Location;
 
-use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::core::types::Range;
+use crate::format::builders::literal;
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 struct FloatAtom {

crates/ruff_python_formatter/src/format/stmt.rs

@@ -4,15 +4,14 @@ use ruff_formatter::prelude::*;
 use ruff_formatter::{format_args, write};
 use ruff_text_size::TextSize;
 
-use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::cst::{
     Alias, Arguments, Excepthandler, Expr, ExprKind, Keyword, Stmt, StmtKind, Withitem,
 };
 use crate::format::builders::{block, join_names};
+use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
 use crate::format::helpers::is_self_closing;
 use crate::shared_traits::AsFormat;
-use crate::trivia::TriviaKind;
 
 fn format_break(f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
     write!(f, [text("break")])
@@ -22,20 +21,7 @@ fn format_pass(f: &mut Formatter<ASTFormatContext<'_>>, stmt: &Stmt) -> FormatRe
     // Write the statement body.
     write!(f, [text("pass")])?;
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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))])?;
-    }
 
     Ok(())
 }
@@ -200,20 +186,7 @@ fn format_func_def(
 
     write!(f, [text(":")])?;
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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, [block_indent(&format_args![block(body)])])
 }
@@ -245,20 +218,7 @@ fn format_assign(
         )?;
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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))])?;
-    }
 
     Ok(())
 }
@@ -417,20 +377,7 @@ fn format_return(
         write!(f, [space(), value.format()])?;
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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))])?;
-    }
 
     Ok(())
 }
@@ -590,20 +537,7 @@ fn format_import_from(
         )?;
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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))])?;
-    }
 
     Ok(())
 }
@@ -635,20 +569,7 @@ fn format_expr(
         )?;
     }
 
-    // Format any end-of-line comments.
+    write!(f, [end_of_line_comments(stmt)])?;
-    let mut first = true;
-    for range in stmt.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))])?;
-    }
 
     Ok(())
 }
@@ -697,20 +618,7 @@ pub struct FormatStmt<'a> {
 
 impl Format<ASTFormatContext<'_>> for FormatStmt<'_> {
     fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
-        // Any leading comments come on the line before.
+        write!(f, [leading_comments(self.item)])?;
-        for trivia in &self.item.trivia {
-            if trivia.relationship.is_leading() {
-                match trivia.kind {
-                    TriviaKind::EmptyLine => {
-                        write!(f, [empty_line()])?;
-                    }
-                    TriviaKind::OwnLineComment(range) => {
-                        write!(f, [literal(range), hard_line_break()])?;
-                    }
-                    _ => {}
-                }
-            }
-        }
 
         match &self.item.node {
             StmtKind::Pass => format_pass(f, self.item),
@@ -852,21 +760,7 @@ impl Format<ASTFormatContext<'_>> for FormatStmt<'_> {
         }?;
 
         write!(f, [hard_line_break()])?;
-
+        write!(f, [trailing_comments(self.item)])?;
-        // Any trailing comments come on the lines after.
-        for trivia in &self.item.trivia {
-            if trivia.relationship.is_trailing() {
-                match trivia.kind {
-                    TriviaKind::EmptyLine => {
-                        write!(f, [empty_line()])?;
-                    }
-                    TriviaKind::OwnLineComment(range) => {
-                        write!(f, [literal(range), hard_line_break()])?;
-                    }
-                    _ => {}
-                }
-            }
-        }
 
         Ok(())
     }

crates/ruff_python_formatter/src/lib.rs

@@ -11,7 +11,6 @@ use crate::newlines::normalize_newlines;
 use crate::parentheses::normalize_parentheses;
 
 mod attachment;
-pub mod builders;
 pub mod cli;
 pub mod context;
 mod core;

crates/ruff_python_formatter/src/main.rs

@@ -2,6 +2,7 @@ use std::fs;
 
 use anyhow::Result;
 use clap::Parser as ClapParser;
+
 use ruff_python_formatter::cli::Cli;
 use ruff_python_formatter::fmt;
 

crates/ruff_python_formatter/src/newlines.rs

@@ -38,13 +38,13 @@ enum Trailer {
     CompoundStatement,
 }
 
-struct NewlineNormalizer {
+struct StmtNormalizer {
     depth: Depth,
     trailer: Trailer,
     scope: Scope,
 }
 
-impl<'a> Visitor<'a> for NewlineNormalizer {
+impl<'a> Visitor<'a> for StmtNormalizer {
     fn visit_stmt(&mut self, stmt: &'a mut Stmt) {
         // Remove any runs of empty lines greater than two in a row.
         let mut count = 0;
@@ -296,12 +296,20 @@ impl<'a> Visitor<'a> for NewlineNormalizer {
         self.depth = prev_depth;
         self.scope = prev_scope;
     }
+}
 
-    fn visit_expr(&mut self, _expr: &'a mut Expr) {}
+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);
+    }
 }
 
 pub fn normalize_newlines(python_cst: &mut [Stmt]) {
-    let mut normalizer = NewlineNormalizer {
+    let mut normalizer = StmtNormalizer {
         depth: Depth::TopLevel,
         trailer: Trailer::None,
         scope: Scope::Module,
@@ -309,4 +317,9 @@ pub fn normalize_newlines(python_cst: &mut [Stmt]) {
     for stmt in python_cst.iter_mut() {
         normalizer.visit_stmt(stmt);
     }
+
+    let mut normalizer = ExprNormalizer;
+    for stmt in python_cst.iter_mut() {
+        normalizer.visit_stmt(stmt);
+    }
 }

crates/ruff_python_formatter/src/parentheses.rs

@@ -1,3 +1,5 @@
+use rustpython_parser::ast::Constant;
+
 use crate::core::helpers::is_radix_literal;
 use crate::core::locator::Locator;
 use crate::core::types::Range;
@@ -5,7 +7,6 @@ use crate::core::visitor;
 use crate::core::visitor::Visitor;
 use crate::cst::{Expr, ExprKind, Stmt, StmtKind};
 use crate::trivia::Parenthesize;
-use rustpython_parser::ast::Constant;
 
 /// Modify an [`Expr`] to infer parentheses, rather than respecting any user-provided trivia.
 fn use_inferred_parens(expr: &mut Expr) {