Use one line between top-level items if formatting a stub file (#6501)

Co-authored-by: Micha Reiser <micha@reiser.io>

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

@@ -0,0 +1,39 @@
+class A:
+    def __init__(self):
+        pass
+
+class B:
+    def __init__(self):
+        pass
+
+def foo():
+    pass
+
+class Del(expr_context): ...
+class Load(expr_context): ...
+
+# Some comment.
+class Other(expr_context): ...
+class Store(expr_context): ...
+class Foo(Bar): ...
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+def bar(): ...
+def baz(): ...
+def quux():
+    """Some docstring."""
+
+def quuz():
+    """Some docstring."""

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

@@ -0,0 +1,78 @@
+class A:
+    def __init__(self):
+        pass
+
+
+class B:
+    def __init__(self):
+        pass
+
+
+def foo():
+    pass
+
+
+class Del(expr_context):
+    ...
+
+
+class Load(expr_context):
+    ...
+
+
+# Some comment.
+class Other(expr_context):
+    ...
+
+
+class Store(expr_context):
+    ...
+
+
+class Foo(Bar):
+    ...
+
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+
+def bar():
+    ...
+
+
+def baz():
+    ...
+
+
+def quux():
+    """Some docstring."""
+
+
+def quuz():
+    """Some docstring."""
+
+def a():
+  ...
+
+class Test:
+  ...
+
+class Test2(A):
+  ...
+
+def b(): ...
+# comment
+def c(): ...

crates/ruff_python_formatter/src/options.rs

@@ -72,6 +72,10 @@ impl PyFormatOptions {
         self.quote_style
     }
 
+    pub fn source_type(&self) -> PySourceType {
+        self.source_type
+    }
+
     #[must_use]
     pub fn with_quote_style(mut self, style: QuoteStyle) -> Self {
         self.quote_style = style;

crates/ruff_python_formatter/src/statement/suite.rs

@@ -2,7 +2,7 @@ use crate::comments::{leading_comments, trailing_comments};
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
 use ruff_python_ast::helpers::is_compound_statement;
 use ruff_python_ast::node::AnyNodeRef;
-use ruff_python_ast::{self as ast, Expr, ExprConstant, Ranged, Stmt, Suite};
+use ruff_python_ast::{self as ast, Constant, Expr, ExprConstant, Ranged, Stmt, Suite};
 use ruff_python_trivia::{lines_after_ignoring_trivia, lines_before};
 use ruff_text_size::TextRange;
 
@@ -55,6 +55,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
 
         let comments = f.context().comments().clone();
         let source = f.context().source();
+        let source_type = f.options().source_type();
 
         let mut f = WithNodeLevel::new(node_level, f);
         write!(
@@ -152,6 +153,44 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                         || is_class_or_function_definition(following)
                     {
                         match self.kind {
+                            SuiteKind::TopLevel if source_type.is_stub() => {
+                                // Preserve the empty line if the definitions are separated by a comment
+                                if comments.has_trailing_comments(preceding)
+                                    || comments.has_leading_comments(following)
+                                {
+                                    empty_line().fmt(f)?;
+                                } else {
+                                    // Two subsequent classes that both have an ellipsis only body
+                                    // ```python
+                                    // class A: ...
+                                    // class B: ...
+                                    // ```
+                                    let class_sequences_with_ellipsis_only =
+                                        preceding.as_class_def_stmt().is_some_and(|class| {
+                                            contains_only_an_ellipsis(&class.body)
+                                        }) && following.as_class_def_stmt().is_some_and(|class| {
+                                            contains_only_an_ellipsis(&class.body)
+                                        });
+
+                                    // Two subsequent functions where the preceding has an ellipsis only body
+                                    // ```python
+                                    // def test(): ...
+                                    // def b(): a
+                                    // ```
+                                    let function_with_ellipsis =
+                                        preceding.as_function_def_stmt().is_some_and(|function| {
+                                            contains_only_an_ellipsis(&function.body)
+                                        }) && following.is_function_def_stmt();
+
+                                    // Don't add an empty line between two classes that have an `...` body only or after
+                                    // a function with an `...` body. Otherwise add an empty line.
+                                    if !class_sequences_with_ellipsis_only
+                                        && !function_with_ellipsis
+                                    {
+                                        empty_line().fmt(f)?;
+                                    }
+                                }
+                            }
                             SuiteKind::TopLevel => {
                                 write!(f, [empty_line(), empty_line()])?;
                             }
@@ -284,6 +323,20 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
     }
 }
 
+/// Returns `true` if a function or class body contains only an ellipsis.
+fn contains_only_an_ellipsis(body: &[Stmt]) -> bool {
+    match body {
+        [Stmt::Expr(ast::StmtExpr { value, .. })] => matches!(
+            value.as_ref(),
+            Expr::Constant(ast::ExprConstant {
+                value: Constant::Ellipsis,
+                ..
+            })
+        ),
+        _ => false,
+    }
+}
+
 /// Returns `true` if a [`Stmt`] is a class or function definition.
 const fn is_class_or_function_definition(stmt: &Stmt) -> bool {
     matches!(stmt, Stmt::FunctionDef(_) | Stmt::ClassDef(_))

crates/ruff_python_formatter/tests/fixtures.rs

@@ -161,7 +161,11 @@ fn format() {
         });
     };
 
-    insta::glob!("../resources", "test/fixtures/ruff/**/*.py", test_file);
+    insta::glob!(
+        "../resources",
+        "test/fixtures/ruff/**/*.{py,pyi}",
+        test_file
+    );
 }
 
 /// Format another time and make sure that there are no changes anymore

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

@@ -0,0 +1,118 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/top_level.py
+---
+## Input
+```py
+class A:
+    def __init__(self):
+        pass
+
+class B:
+    def __init__(self):
+        pass
+
+def foo():
+    pass
+
+class Del(expr_context): ...
+class Load(expr_context): ...
+
+# Some comment.
+class Other(expr_context): ...
+class Store(expr_context): ...
+class Foo(Bar): ...
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+def bar(): ...
+def baz(): ...
+def quux():
+    """Some docstring."""
+
+def quuz():
+    """Some docstring."""
+```
+
+## Output
+```py
+class A:
+    def __init__(self):
+        pass
+
+
+class B:
+    def __init__(self):
+        pass
+
+
+def foo():
+    pass
+
+
+class Del(expr_context):
+    ...
+
+
+class Load(expr_context):
+    ...
+
+
+# Some comment.
+class Other(expr_context):
+    ...
+
+
+class Store(expr_context):
+    ...
+
+
+class Foo(Bar):
+    ...
+
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+
+def bar():
+    ...
+
+
+def baz():
+    ...
+
+
+def quux():
+    """Some docstring."""
+
+
+def quuz():
+    """Some docstring."""
+```
+
+
+

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

@@ -0,0 +1,153 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/top_level.pyi
+---
+## Input
+```py
+class A:
+    def __init__(self):
+        pass
+
+
+class B:
+    def __init__(self):
+        pass
+
+
+def foo():
+    pass
+
+
+class Del(expr_context):
+    ...
+
+
+class Load(expr_context):
+    ...
+
+
+# Some comment.
+class Other(expr_context):
+    ...
+
+
+class Store(expr_context):
+    ...
+
+
+class Foo(Bar):
+    ...
+
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+
+def bar():
+    ...
+
+
+def baz():
+    ...
+
+
+def quux():
+    """Some docstring."""
+
+
+def quuz():
+    """Some docstring."""
+
+def a():
+  ...
+
+class Test:
+  ...
+
+class Test2(A):
+  ...
+
+def b(): ...
+# comment
+def c(): ...
+```
+
+## Output
+```py
+class A:
+    def __init__(self):
+        pass
+
+class B:
+    def __init__(self):
+        pass
+
+def foo():
+    pass
+
+class Del(expr_context):
+    ...
+class Load(expr_context):
+    ...
+
+# Some comment.
+class Other(expr_context):
+    ...
+class Store(expr_context):
+    ...
+class Foo(Bar):
+    ...
+
+class Baz(Qux):
+    def __init__(self):
+        pass
+
+class Quux(Qux):
+    def __init__(self):
+        pass
+
+# Some comment.
+class Quuz(Qux):
+    def __init__(self):
+        pass
+
+def bar():
+    ...
+def baz():
+    ...
+def quux():
+    """Some docstring."""
+
+def quuz():
+    """Some docstring."""
+
+def a():
+    ...
+
+class Test:
+    ...
+class Test2(A):
+    ...
+
+def b():
+    ...
+
+# comment
+def c():
+    ...
+```
+
+
+