Format function and class definitions into a single line if its body is an ellipsis (#6592)

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/ellipsis.pyi

@@ -0,0 +1,102 @@
+"""Compound statements with no body should be written on one line."""
+
+if True:
+    ...
+elif True:
+    ...
+else:
+    ...
+
+if True:
+    # comment
+    ...
+elif True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+if True:
+    ... # comment
+elif True:
+    ... # comment
+else:
+    ... # comment
+
+for i in []:
+    ...
+else:
+    ...
+
+for i in []:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+for i in []:
+    ... # comment
+else:
+    ... # comment
+
+while True:
+    ...
+else:
+    ...
+
+while True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+while True:
+    ... # comment
+else:
+    ... # comment
+
+with True:
+    ...
+
+with True:
+    # comment
+    ...
+
+with True:
+    ... # comment
+
+match x:
+    case 1:
+        ...
+    case 2:
+        # comment
+        ...
+    case 3:
+        ... # comment
+
+try:
+    ...
+except:
+    ...
+finally:
+    ...
+
+try:
+    # comment
+    ...
+except:
+    # comment
+    ...
+finally:
+    # comment
+    ...
+
+try:
+    ...  # comment
+except:
+    ...  # comment
+finally:
+    ...  # comment

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/top_level.pyi

@@ -76,3 +76,16 @@ class Test2(A):
 def b(): ...
 # comment
 def c(): ...
+
+class EllipsisWithComment:
+    ...  # comment
+
+def function_with_comment():
+    ...  # comment
+
+class EllispsisWithMultipleTrailing: # trailing class comment
+    ... # trailing ellipsis comment
+
+class EllipsisWithLeadingComment:
+    # leading
+    ...

crates/ruff_python_formatter/src/other/except_handler_except_handler.rs

@@ -6,7 +6,7 @@ use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::{FormatNodeRule, PyFormatter};
 
 #[derive(Copy, Clone, Default)]
@@ -86,7 +86,7 @@ impl FormatNodeRule<ExceptHandlerExceptHandler> for FormatExceptHandlerExceptHan
                         Ok(())
                     }),
                 ),
-                block_indent(&body.format())
+                clause_body(body, dangling_comments),
             ]
         )
     }

crates/ruff_python_formatter/src/other/match_case.rs

@@ -6,7 +6,7 @@ use ruff_text_size::TextRange;
 use crate::comments::{leading_comments, SourceComment};
 use crate::expression::parentheses::parenthesized;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::{FormatError, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -57,7 +57,7 @@ impl FormatNodeRule<MatchCase> for FormatMatchCase {
                         Ok(())
                     }),
                 ),
-                block_indent(&body.format())
+                clause_body(body, dangling_item_comments),
             ]
         )
     }

crates/ruff_python_formatter/src/statement/clause.rs

@@ -2,12 +2,13 @@ use crate::comments::{
     leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind,
 };
 use crate::prelude::*;
+use crate::statement::suite::{contains_only_an_ellipsis, SuiteKind};
 use crate::verbatim::write_suppressed_clause_header;
-use ruff_formatter::{Argument, Arguments, FormatError};
+use ruff_formatter::{write, Argument, Arguments, FormatError};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{
     ElifElseClause, ExceptHandlerExceptHandler, MatchCase, Ranged, StmtClassDef, StmtFor,
-    StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith,
+    StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith, Suite,
 };
 use ruff_python_trivia::{SimpleToken, SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::{TextRange, TextSize};
@@ -352,6 +353,57 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatClauseHeader<'_, 'ast> {
     }
 }
 
+pub(crate) struct FormatClauseBody<'a> {
+    body: &'a Suite,
+    kind: SuiteKind,
+    trailing_comments: &'a [SourceComment],
+}
+
+impl<'a> FormatClauseBody<'a> {
+    #[must_use]
+    pub(crate) fn with_kind(mut self, kind: SuiteKind) -> Self {
+        self.kind = kind;
+        self
+    }
+}
+
+pub(crate) fn clause_body<'a>(
+    body: &'a Suite,
+    trailing_comments: &'a [SourceComment],
+) -> FormatClauseBody<'a> {
+    FormatClauseBody {
+        body,
+        kind: SuiteKind::default(),
+        trailing_comments,
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatClauseBody<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        if f.options().source_type().is_stub()
+            && contains_only_an_ellipsis(self.body, f.context().comments())
+            && self.trailing_comments.is_empty()
+        {
+            write!(
+                f,
+                [
+                    space(),
+                    self.body.format().with_options(self.kind),
+                    hard_line_break()
+                ]
+            )
+        } else {
+            write!(
+                f,
+                [
+                    trailing_comments(self.trailing_comments),
+                    block_indent(&self.body.format().with_options(self.kind))
+                ]
+            )
+        }
+    }
+}
+
 /// Finds the range of `keyword` starting the search at `start_position`. Expects only comments and `(` between
 /// the `start_position` and the `keyword` token.
 fn find_keyword(

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -6,7 +6,7 @@ use crate::comments::{leading_comments, trailing_comments, SourceComment};
 use crate::prelude::*;
 use crate::statement::suite::SuiteKind;
 
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::FormatNodeRule;
 
 #[derive(Default)]
@@ -107,7 +107,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
                         Ok(())
                     }),
                 ),
-                block_indent(&body.format().with_options(SuiteKind::Class))
+                clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Class),
             ]
         )
     }

crates/ruff_python_formatter/src/statement/stmt_for.rs

@@ -6,7 +6,7 @@ use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader, ElseClause};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
 use crate::FormatNodeRule;
 
 #[derive(Debug)]
@@ -64,7 +64,7 @@ impl FormatNodeRule<StmtFor> for FormatStmtFor {
                         maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks),
                     ],
                 ),
-                block_indent(&body.format())
+                clause_body(body, trailing_condition_comments),
             ]
         )?;
 
@@ -86,7 +86,7 @@ impl FormatNodeRule<StmtFor> for FormatStmtFor {
                         &text("else"),
                     )
                     .with_leading_comments(leading, body.last()),
-                    block_indent(&orelse.format())
+                    clause_body(orelse, trailing),
                 ]
             )?;
         }

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -6,7 +6,7 @@ use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{Parentheses, Parenthesize};
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::statement::stmt_class_def::FormatDecorators;
 use crate::statement::suite::SuiteKind;
 use crate::FormatNodeRule;
@@ -142,7 +142,7 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
                         group(&format_inner).fmt(f)
                     }),
                 ),
-                block_indent(&body.format().with_options(SuiteKind::Function))
+                clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Function),
             ]
         )
     }

crates/ruff_python_formatter/src/statement/stmt_if.rs

@@ -6,7 +6,7 @@ use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::FormatNodeRule;
 
 #[derive(Default)]
@@ -36,7 +36,7 @@ impl FormatNodeRule<StmtIf> for FormatStmtIf {
                         maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
                     ],
                 ),
-                block_indent(&body.format())
+                clause_body(body, trailing_colon_comment),
             ]
         )?;
 
@@ -98,7 +98,7 @@ pub(crate) fn format_elif_else_clause(
                 }),
             )
             .with_leading_comments(leading_comments, last_node),
-            block_indent(&body.format())
+            clause_body(body, trailing_colon_comment),
         ]
     )
 }

crates/ruff_python_formatter/src/statement/stmt_try.rs

@@ -6,7 +6,7 @@ use crate::comments::leading_alternate_branch_comments;
 use crate::comments::SourceComment;
 use crate::other::except_handler_except_handler::ExceptHandlerKind;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader, ElseClause};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
 use crate::statement::{FormatRefWithRule, Stmt};
 use crate::{FormatNodeRule, PyFormatter};
 
@@ -138,7 +138,7 @@ fn format_case<'a>(
             [
                 clause_header(header, trailing_case_comments, &text(kind.keyword()))
                     .with_leading_comments(leading_case_comments, previous_node),
-                block_indent(&body.format())
+                clause_body(body, trailing_case_comments),
             ]
         )?;
         (None, rest)

crates/ruff_python_formatter/src/statement/stmt_while.rs

@@ -6,7 +6,7 @@ use crate::comments::SourceComment;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader, ElseClause};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
 use crate::FormatNodeRule;
 
 #[derive(Default)]
@@ -43,7 +43,7 @@ impl FormatNodeRule<StmtWhile> for FormatStmtWhile {
                         maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
                     ]
                 ),
-                block_indent(&body.format())
+                clause_body(body, trailing_condition_comments),
             ]
         )?;
 
@@ -63,7 +63,7 @@ impl FormatNodeRule<StmtWhile> for FormatStmtWhile {
                         &text("else")
                     )
                     .with_leading_comments(leading, body.last()),
-                    block_indent(&orelse.format())
+                    clause_body(orelse, trailing),
                 ]
             )?;
         }

crates/ruff_python_formatter/src/statement/stmt_with.rs

@@ -11,7 +11,7 @@ use crate::expression::parentheses::{
 };
 use crate::other::commas;
 use crate::prelude::*;
-use crate::statement::clause::{clause_header, ClauseHeader};
+use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::{FormatNodeRule, PyFormatOptions};
 
 #[derive(Default)]
@@ -100,7 +100,7 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
                         Ok(())
                     })
                 ),
-                block_indent(&item.body.format())
+                clause_body(&item.body, colon_comments)
             ]
         )
     }

crates/ruff_python_formatter/src/statement/suite.rs

@@ -1,4 +1,4 @@
-use crate::comments::{leading_comments, trailing_comments};
+use crate::comments::{leading_comments, trailing_comments, Comments};
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
 use ruff_python_ast::helpers::is_compound_statement;
 use ruff_python_ast::node::AnyNodeRef;
@@ -17,7 +17,7 @@ use crate::verbatim::{
 };
 
 /// Level at which the [`Suite`] appears in the source code.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, Default)]
 pub enum SuiteKind {
     /// Statements at the module level / top level
     TopLevel,
@@ -29,6 +29,7 @@ pub enum SuiteKind {
     Class,
 
     /// Statements in any other body (e.g., `if` or `while`).
+    #[default]
     Other,
 }
 
@@ -168,9 +169,15 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                                     // ```
                                     let class_sequences_with_ellipsis_only =
                                         preceding.as_class_def_stmt().is_some_and(|class| {
-                                            contains_only_an_ellipsis(&class.body)
+                                            contains_only_an_ellipsis(
+                                                &class.body,
+                                                f.context().comments(),
+                                            )
                                         }) && following.as_class_def_stmt().is_some_and(|class| {
-                                            contains_only_an_ellipsis(&class.body)
+                                            contains_only_an_ellipsis(
+                                                &class.body,
+                                                f.context().comments(),
+                                            )
                                         });
 
                                     // Two subsequent functions where the preceding has an ellipsis only body
@@ -180,7 +187,10 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                                     // ```
                                     let function_with_ellipsis =
                                         preceding.as_function_def_stmt().is_some_and(|function| {
-                                            contains_only_an_ellipsis(&function.body)
+                                            contains_only_an_ellipsis(
+                                                &function.body,
+                                                f.context().comments(),
+                                            )
                                         }) && following.is_function_def_stmt();
 
                                     // Don't add an empty line between two classes that have an `...` body only or after
@@ -325,16 +335,19 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
     }
 }
 
-/// Returns `true` if a function or class body contains only an ellipsis.
+/// Returns `true` if a function or class body contains only an ellipsis with no comments.
-fn contains_only_an_ellipsis(body: &[Stmt]) -> bool {
+pub(crate) fn contains_only_an_ellipsis(body: &[Stmt], comments: &Comments) -> bool {
     match body {
-        [Stmt::Expr(ast::StmtExpr { value, .. })] => matches!(
+        [Stmt::Expr(ast::StmtExpr { value, .. })] => {
+            let [node] = body else { return false; };
+            matches!(
                 value.as_ref(),
                 Expr::Constant(ast::ExprConstant {
                     value: Constant::Ellipsis,
                     ..
                 })
-        ),
+            ) && !comments.has_leading(node)
+        }
         _ => false,
     }
 }

crates/ruff_python_formatter/tests/snapshots/format@statement__ellipsis.pyi.snap

@@ -0,0 +1,193 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/ellipsis.pyi
+---
+## Input
+```py
+"""Compound statements with no body should be written on one line."""
+
+if True:
+    ...
+elif True:
+    ...
+else:
+    ...
+
+if True:
+    # comment
+    ...
+elif True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+if True:
+    ... # comment
+elif True:
+    ... # comment
+else:
+    ... # comment
+
+for i in []:
+    ...
+else:
+    ...
+
+for i in []:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+for i in []:
+    ... # comment
+else:
+    ... # comment
+
+while True:
+    ...
+else:
+    ...
+
+while True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+while True:
+    ... # comment
+else:
+    ... # comment
+
+with True:
+    ...
+
+with True:
+    # comment
+    ...
+
+with True:
+    ... # comment
+
+match x:
+    case 1:
+        ...
+    case 2:
+        # comment
+        ...
+    case 3:
+        ... # comment
+
+try:
+    ...
+except:
+    ...
+finally:
+    ...
+
+try:
+    # comment
+    ...
+except:
+    # comment
+    ...
+finally:
+    # comment
+    ...
+
+try:
+    ...  # comment
+except:
+    ...  # comment
+finally:
+    ...  # comment```
+
+## Output
+```py
+"""Compound statements with no body should be written on one line."""
+
+if True: ...
+elif True: ...
+else: ...
+
+if True:
+    # comment
+    ...
+elif True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+if True: ...  # comment
+elif True: ...  # comment
+else: ...  # comment
+
+for i in []: ...
+else: ...
+
+for i in []:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+for i in []: ...  # comment
+else: ...  # comment
+
+while True: ...
+else: ...
+
+while True:
+    # comment
+    ...
+else:
+    # comment
+    ...
+
+while True: ...  # comment
+else: ...  # comment
+
+with True: ...
+
+with True:
+    # comment
+    ...
+
+with True: ...  # comment
+
+match x:
+    case "NOT_YET_IMPLEMENTED_PatternMatchValue": ...
+    case "NOT_YET_IMPLEMENTED_PatternMatchValue":
+        # comment
+        ...
+    case "NOT_YET_IMPLEMENTED_PatternMatchValue": ...  # comment
+
+try: ...
+except: ...
+finally: ...
+
+try:
+    # comment
+    ...
+except:
+    # comment
+    ...
+finally:
+    # comment
+    ...
+
+try: ...  # comment
+except: ...  # comment
+finally: ...  # comment
+```
+
+
+

crates/ruff_python_formatter/tests/snapshots/format@statement__top_level.pyi.snap

@@ -82,6 +82,19 @@ class Test2(A):
 def b(): ...
 # comment
 def c(): ...
+
+class EllipsisWithComment:
+    ...  # comment
+
+def function_with_comment():
+    ...  # comment
+
+class EllispsisWithMultipleTrailing: # trailing class comment
+    ... # trailing ellipsis comment
+
+class EllipsisWithLeadingComment:
+    # leading
+    ...
 ```
 
 ## Output
@@ -97,18 +110,13 @@ class B:
 def foo():
     pass
 
-class Del(expr_context):
+class Del(expr_context): ...
-    ...
+class Load(expr_context): ...
-class Load(expr_context):
-    ...
 
 # Some comment.
-class Other(expr_context):
+class Other(expr_context): ...
-    ...
+class Store(expr_context): ...
-class Store(expr_context):
+class Foo(Bar): ...
-    ...
-class Foo(Bar):
-    ...
 
 class Baz(Qux):
     def __init__(self):
@@ -123,29 +131,33 @@ class Quuz(Qux):
     def __init__(self):
         pass
 
-def bar():
+def bar(): ...
-    ...
+def baz(): ...
-def baz():
-    ...
 def quux():
     """Some docstring."""
 
 def quuz():
     """Some docstring."""
 
-def a():
+def a(): ...
-    ...
 
-class Test:
+class Test: ...
-    ...
+class Test2(A): ...
-class Test2(A):
-    ...
 
-def b():
+def b(): ...
-    ...
 
 # comment
-def c():
+def c(): ...
+
+class EllipsisWithComment: ...  # comment
+
+def function_with_comment(): ...  # comment
+
+class EllispsisWithMultipleTrailing:  # trailing class comment
+    ...  # trailing ellipsis comment
+
+class EllipsisWithLeadingComment:
+    # leading
     ...
 ```