Leading, Dangling, and Trailing comments formatting

crates/ruff_formatter/src/source_code.rs

@@ -1,4 +1,4 @@
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextRange, TextSize};
 use std::fmt::{Debug, Formatter};
 
 /// The source code of a document that gets formatted
@@ -68,9 +68,17 @@ impl SourceCodeSlice {
         &code.text[self.range]
     }
 
-    pub fn range(&self) -> TextRange {
+    pub const fn range(&self) -> TextRange {
         self.range
     }
+
+    pub const fn start(&self) -> TextSize {
+        self.range.start()
+    }
+
+    pub const fn end(&self) -> TextSize {
+        self.range.end()
+    }
 }
 
 impl Debug for SourceCodeSlice {

crates/ruff_python_formatter/src/comments/format.rs

@@ -0,0 +1,182 @@
+use crate::context::NodeLevel;
+use crate::prelude::*;
+use crate::trivia::{lines_after, lines_before};
+use ruff_formatter::{format_args, write};
+use ruff_python_ast::node::AnyNodeRef;
+use ruff_python_ast::prelude::AstNode;
+
+/// Formats the leading comments of a node.
+pub(crate) fn leading_comments<T>(node: &T) -> FormatLeadingComments
+where
+    T: AstNode,
+{
+    FormatLeadingComments {
+        node: node.as_any_node_ref(),
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub(crate) struct FormatLeadingComments<'a> {
+    node: AnyNodeRef<'a>,
+}
+
+impl Format<PyFormatContext<'_>> for FormatLeadingComments<'_> {
+    fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
+        let comments = f.context().comments().clone();
+
+        for comment in comments.leading_comments(self.node) {
+            let slice = comment.slice();
+
+            let lines_after_comment = lines_after(f.context().contents(), slice.end()).max(1);
+            write!(
+                f,
+                [
+                    source_text_slice(slice.range(), ContainsNewlines::No),
+                    empty_lines(lines_after_comment)
+                ]
+            )?;
+
+            comment.mark_formatted();
+        }
+
+        Ok(())
+    }
+}
+
+/// Formats the trailing comments of `node`
+pub(crate) fn trailing_comments<T>(node: &T) -> FormatTrailingComments
+where
+    T: AstNode,
+{
+    FormatTrailingComments {
+        node: node.as_any_node_ref(),
+    }
+}
+
+pub(crate) struct FormatTrailingComments<'a> {
+    node: AnyNodeRef<'a>,
+}
+
+impl Format<PyFormatContext<'_>> for FormatTrailingComments<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let comments = f.context().comments().clone();
+        let mut has_empty_lines_before = false;
+
+        for trailing in comments.trailing_comments(self.node) {
+            let slice = trailing.slice();
+            let content = source_text_slice(slice.range(), ContainsNewlines::No);
+
+            let lines_before_comment = lines_before(f.context().contents(), slice.start());
+            has_empty_lines_before |= lines_before_comment > 0;
+
+            if has_empty_lines_before {
+                // A trailing comment at the end of a body or list
+                // ```python
+                // def test():
+                //      pass
+                //
+                //      # Some comment
+                // ```
+                write!(
+                    f,
+                    [
+                        line_suffix(&format_with(|f| {
+                            write!(f, [empty_lines(lines_before_comment), content])
+                        })),
+                        expand_parent()
+                    ]
+                )?;
+            } else {
+                write!(
+                    f,
+                    [
+                        line_suffix(&format_args![space(), space(), content]),
+                        expand_parent()
+                    ]
+                )?;
+            }
+
+            trailing.mark_formatted();
+        }
+
+        Ok(())
+    }
+}
+
+/// Formats the dangling comments of `node`.
+pub(crate) fn dangling_comments<T>(node: &T) -> FormatDanglingComments
+where
+    T: AstNode,
+{
+    FormatDanglingComments {
+        node: node.as_any_node_ref(),
+    }
+}
+
+pub(crate) struct FormatDanglingComments<'a> {
+    node: AnyNodeRef<'a>,
+}
+
+impl Format<PyFormatContext<'_>> for FormatDanglingComments<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
+        let comments = f.context().comments().clone();
+
+        let dangling_comments = comments.dangling_comments(self.node);
+
+        let mut first = true;
+        for comment in dangling_comments {
+            if first && comment.position().is_end_of_line() {
+                write!(f, [space(), space()])?;
+            }
+
+            write!(
+                f,
+                [
+                    source_text_slice(comment.slice().range(), ContainsNewlines::No),
+                    empty_lines(lines_after(f.context().contents(), comment.slice().end()))
+                ]
+            )?;
+
+            comment.mark_formatted();
+
+            first = false;
+        }
+
+        Ok(())
+    }
+}
+
+// Helper that inserts the appropriate number of empty lines before a comment, depending on the node level.
+// Top level: Up to two empty lines
+// parenthesized: A single empty line
+// other: Up to a single empty line
+const fn empty_lines(lines: u32) -> FormatEmptyLines {
+    FormatEmptyLines { lines }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct FormatEmptyLines {
+    lines: u32,
+}
+
+impl Format<PyFormatContext<'_>> for FormatEmptyLines {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
+        match f.context().node_level() {
+            NodeLevel::TopLevel | NodeLevel::Statement => match self.lines {
+                0 => Ok(()),
+                1 => write!(f, [hard_line_break()]),
+                lines => {
+                    write!(f, [empty_line()])?;
+                    if lines > 2 && f.context().node_level().is_top_level() {
+                        write!(f, [empty_line()])?
+                    }
+
+                    Ok(())
+                }
+            },
+
+            // Remove all whitespace in parenthesized expressions
+            NodeLevel::Parenthesized => write!(f, [hard_line_break()]),
+        }
+    }
+}

crates/ruff_python_formatter/src/comments/mod.rs

@@ -93,6 +93,7 @@ use std::fmt::Debug;
 use std::rc::Rc;
 
 mod debug;
+mod format;
 mod map;
 mod node_key;
 mod placement;
@@ -102,6 +103,7 @@ use crate::comments::debug::{DebugComment, DebugComments};
 use crate::comments::map::MultiMap;
 use crate::comments::node_key::NodeRefEqualityKey;
 use crate::comments::visitor::CommentsVisitor;
+pub(crate) use format::{dangling_comments, leading_comments, trailing_comments};
 use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::source_code::CommentRanges;

crates/ruff_python_formatter/src/context.rs

@@ -8,6 +8,7 @@ pub struct PyFormatContext<'a> {
     options: SimpleFormatOptions,
     contents: &'a str,
     comments: Comments<'a>,
+    node_level: NodeLevel,
 }
 
 impl<'a> PyFormatContext<'a> {
@@ -20,6 +21,7 @@ impl<'a> PyFormatContext<'a> {
             options,
             contents,
             comments,
+            node_level: NodeLevel::TopLevel,
         }
     }
 
@@ -33,6 +35,15 @@ impl<'a> PyFormatContext<'a> {
         Locator::new(self.contents)
     }
 
+    #[allow(unused)]
+    pub(crate) fn set_node_level(&mut self, level: NodeLevel) {
+        self.node_level = level;
+    }
+
+    pub(crate) fn node_level(&self) -> NodeLevel {
+        self.node_level
+    }
+
     #[allow(unused)]
     pub(crate) fn comments(&self) -> &Comments<'a> {
         &self.comments
@@ -56,7 +67,30 @@ impl Debug for PyFormatContext<'_> {
         f.debug_struct("PyFormatContext")
             .field("options", &self.options)
             .field("comments", &self.comments.debug(self.source_code()))
+            .field("node_level", &self.node_level)
             .field("source", &self.contents)
             .finish()
     }
 }
+
+/// What's the enclosing level of the outer node.
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
+pub(crate) enum NodeLevel {
+    /// Formatting statements on the module level.
+    #[default]
+    TopLevel,
+
+    /// Formatting nodes that are enclosed by a statement.
+    #[allow(unused)]
+    Statement,
+
+    /// Formatting nodes that are enclosed in a parenthesized expression.
+    #[allow(unused)]
+    Parenthesized,
+}
+
+impl NodeLevel {
+    pub(crate) const fn is_top_level(self) -> bool {
+        matches!(self, Self::TopLevel)
+    }
+}

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,7 @@ use ruff_formatter::{
 use ruff_python_ast::node::AstNode;
 use ruff_python_ast::source_code::{CommentRanges, CommentRangesBuilder, Locator};
 
-use crate::comments::Comments;
+use crate::comments::{dangling_comments, leading_comments, trailing_comments, Comments};
 use crate::context::PyFormatContext;
 
 pub mod cli;
@@ -57,31 +57,31 @@ where
     /// Formats the node's fields.
     fn fmt_fields(&self, item: &N, f: &mut PyFormatter) -> FormatResult<()>;
 
-    /// Formats the [leading comments](crate::comments#leading-comments) of the node.
+    /// Formats the [leading comments](comments#leading-comments) of the node.
     ///
     /// You may want to override this method if you want to manually handle the formatting of comments
     /// inside of the `fmt_fields` method or customize the formatting of the leading comments.
-    fn fmt_leading_comments(&self, _node: &N, _f: &mut PyFormatter) -> FormatResult<()> {
-        Ok(())
+    fn fmt_leading_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
+        leading_comments(node).fmt(f)
     }
 
-    /// Formats the [dangling comments](rome_formatter::comments#dangling-comments) of the node.
+    /// Formats the [dangling comments](comments#dangling-comments) of the node.
     ///
     /// You should override this method if the node handled by this rule can have dangling comments because the
     /// default implementation formats the dangling comments at the end of the node, which isn't ideal but ensures that
     /// no comments are dropped.
     ///
     /// A node can have dangling comments if all its children are tokens or if all node childrens are optional.
-    fn fmt_dangling_comments(&self, _node: &N, _f: &mut PyFormatter) -> FormatResult<()> {
-        Ok(())
+    fn fmt_dangling_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
+        dangling_comments(node).fmt(f)
     }
 
-    /// Formats the [trailing comments](rome_formatter::comments#trailing-comments) of the node.
+    /// Formats the [trailing comments](comments#trailing-comments) of the node.
     ///
     /// You may want to override this method if you want to manually handle the formatting of comments
     /// inside of the `fmt_fields` method or customize the formatting of the trailing comments.
-    fn fmt_trailing_comments(&self, _node: &N, _f: &mut PyFormatter) -> FormatResult<()> {
-        Ok(())
+    fn fmt_trailing_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
+        trailing_comments(node).fmt(f)
     }
 }
 
@@ -180,8 +180,10 @@ if    True:
     print( "hi" )
 # trailing
 "#;
-        let expected = r#"if    True:
+        let expected = r#"# preceding
+if    True:
     print( "hi" )
+# trailing
 "#;
         let actual = format_module(input)?.as_code().to_string();
         assert_eq!(expected, actual);