Format empty lines in stub files like black's preview style

.editorconfig

@@ -10,7 +10,7 @@ indent_style = space
 insert_final_newline = true
 indent_size = 2
 
-[*.{rs,py}]
+[*.{rs,py,pyi}]
 indent_size = 4
 
 [*.snap]

crates/ruff_python_formatter/resources/test/fixtures/ruff/module_dangling_comment1.py

@@ -0,0 +1,3 @@
+__all__ = ["X", "XK", "Xatom", "Xcursorfont", "Xutil", "display", "error", "rdb"]
+
+# Shared types throughout the stub

crates/ruff_python_formatter/resources/test/fixtures/ruff/module_dangling_comment2.py

@@ -0,0 +1,2 @@
+__all__ = ["X", "XK", "Xatom", "Xcursorfont", "Xutil", "display", "error", "rdb"]
+# Shared types throughout the stub

crates/ruff_python_formatter/resources/test/fixtures/ruff/stub_files/comments.pyi

@@ -0,0 +1,7 @@
+class SupportsAnext:
+    def __anext__(self): ...
+
+# Comparison protocols
+
+class SupportsDunderLT:
+    def __init__(self): ...

crates/ruff_python_formatter/resources/test/fixtures/ruff/stub_files/nesting.pyi

@@ -0,0 +1,31 @@
+"""Tests specifically for https://github.com/psf/black/issues/3861"""
+
+import sys
+
+
+class OuterClassOrOtherSuite:
+    class Nested11:
+        class Nested12:
+            assignment = 1
+            def function_definition(self): ...
+
+    def f1(self) -> str: ...
+
+    class Nested21:
+        class Nested22:
+            def function_definition(self): ...
+            assignment = 1
+
+    def f2(self) -> str: ...
+
+if sys.version_info > (3, 7):
+    if sys.platform == "win32":
+        assignment = 1
+        def function_definition(self): ...
+
+    def f1(self) -> str: ...
+    if sys.platform != "win32":
+        def function_definition(self): ...
+        assignment = 1
+
+    def f2(self) -> str: ...

crates/ruff_python_formatter/resources/test/fixtures/ruff/stub_files/suite.pyi

@@ -0,0 +1,149 @@
+"""Tests for empty line rules in stub files, mostly inspired by typeshed.
+The rules are a list of nested exceptions. See also
+https://github.com/psf/black/blob/c160e4b7ce30c661ac4f2dfa5038becf1b8c5c33/src/black/lines.py#L576-L744
+"""
+
+import sys
+from typing import Self, TypeAlias, final
+
+if sys.version_info >= (3, 8):
+    class InnerClass1: ...
+
+    class InnerClass2:
+        def a(self): ...
+
+    class InnerClass3:
+        def a(self): ...
+
+    class InnerClass4: ...
+    details: int
+    def f1(self, hresult: int, text: str | None, detail: int) -> None: ...
+    details: int
+    def f2(self, hresult: int, text: str | None, detail: int) -> None: ...
+    @final
+    class DecoratorInsteadOfEmptyLine: ...
+
+    def open(device: str) -> None: ...
+
+    # oss_mixer_device return type
+    def openmixer(device: str = ...) -> None: ...
+    def open2(device: str) -> None: ...
+    # oss_mixer_device2 return type
+    def openmixer2(device: str = ...) -> None: ...
+
+else:
+    class Slice1: ...
+    _Slice1: TypeAlias = Slice1
+
+    class Slice2: ...
+    _Slice2: TypeAlias = Slice2
+
+class NoEmptyLinesBetweenFunctions:
+    def multi_line_but_only_ellipsis(
+        self,
+        mandatory_release: float | None,
+    ) -> None: ...
+    def only_ellipsis1(self) -> float: ...
+    def only_ellipsis2(self) -> float | None: ...
+    def has_impl1(self):
+        print(self)
+        return 1
+
+    def has_impl2(self):
+        print(self)
+        return 2
+
+    def no_impl4(self): ...
+
+class NoEmptyLinesBetweenField:
+    field1: int
+    field2: (
+        # type
+        int
+    )
+    field3 = 3
+    field4 = (
+        1,
+        2,
+    )
+    field5 = 5
+
+class FieldAndFunctionsWithOptionalEmptyLines:
+    details1: int
+    def f1(self, hresult: int, text: str | None, detail: int) -> None: ...
+    details2: int
+    def f2(self, hresult: int, text: str | None, detail: int) -> None: ...
+    details3: int
+
+class NewlinesBetweenStubInnerClasses:
+    def f1(self): ...
+
+    class InnerClass1: ...
+    class InnerClass2: ...
+
+    def f2(self): ...
+
+    class InnerClass3: ...
+    class InnerClass4: ...
+    field = 1
+
+    class InnerClass3: ...
+    class InnerClass4: ...
+
+    def f3(self): ...
+    @final
+    class DecoratorInsteadOfEmptyLine: ...
+
+    @final
+    class DecoratorStillEmptyLine: ...
+
+class NewlinesBetweenInnerClasses:
+    class InnerClass1: ...
+
+    class InnerClass2:
+        def a(self): ...
+
+    class InnerClass3:
+        def a(self): ...
+
+    class InnerClass4: ...
+
+    class InnerClass5:
+        def a(self): ...
+    field1 = 1
+
+    class InnerClass6:
+        def a(self): ...
+
+    def f1(self): ...
+
+    class InnerClass7:
+        def a(self): ...
+    print("hi")
+
+    class InnerClass8:
+        def a(self): ...
+
+class ComplexStatements:
+    # didn't match the name in the C implementation,
+    # meaning it is only *safe* to pass it as a keyword argument on 3.12+
+    if sys.version_info >= (3, 12):
+        @classmethod
+        def fromtimestamp(cls, timestamp: float, tz: float | None = ...) -> Self: ...
+    else:
+        @classmethod
+        def fromtimestamp(cls, __timestamp: float, tz: float | None = ...) -> Self: ...
+
+    @classmethod
+    def utcfromtimestamp(cls, __t: float) -> Self: ...
+    if sys.version_info >= (3, 8):
+        @classmethod
+        def now(cls, tz: float | None = None) -> Self: ...
+    else:
+        @classmethod
+        def now(cls, tz: None = None) -> Self: ...
+        @classmethod
+        def now2(cls, tz: float) -> Self: ...
+
+    @classmethod
+    def utcnow(cls) -> Self: ...

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

@@ -0,0 +1,17 @@
+from typing import final
+
+
+def count1(): ...
+def count2(): ...
+@final
+def count3(): ...
+@final
+class LockType1: ...
+
+def count4(): ...
+
+class LockType2: ...
+class LockType3: ...
+
+@final
+class LockType4: ...

crates/ruff_python_formatter/src/comments/format.rs

@@ -2,6 +2,7 @@ use std::borrow::Cow;
 
 use ruff_formatter::{format_args, write, FormatError, FormatOptions, SourceCode};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
+use ruff_python_ast::PySourceType;
 use ruff_python_trivia::{lines_after, lines_after_ignoring_trivia, lines_before};
 use ruff_text_size::{Ranged, TextLen, TextRange};
 
@@ -485,19 +486,30 @@ fn strip_comment_prefix(comment_text: &str) -> FormatResult<&str> {
 ///
 /// This builder will insert two empty lines before the comment.
 /// ```
-pub(crate) const fn empty_lines_before_trailing_comments(
-    comments: &[SourceComment],
-    expected: u32,
-) -> FormatEmptyLinesBeforeTrailingComments {
-    FormatEmptyLinesBeforeTrailingComments { comments, expected }
+pub(crate) fn empty_lines_before_trailing_comments<'a>(
+    f: &PyFormatter,
+    comments: &'a [SourceComment],
+) -> FormatEmptyLinesBeforeTrailingComments<'a> {
+    // Black has different rules for stub vs. non-stub and top level vs. indented
+    let empty_lines = match (f.options().source_type(), f.context().node_level()) {
+        (PySourceType::Stub, NodeLevel::TopLevel) => 1,
+        (PySourceType::Stub, _) => 0,
+        (_, NodeLevel::TopLevel) => 2,
+        (_, _) => 1,
+    };
+
+    FormatEmptyLinesBeforeTrailingComments {
+        comments,
+        empty_lines,
+    }
 }
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct FormatEmptyLinesBeforeTrailingComments<'a> {
     /// The trailing comments of the node.
     comments: &'a [SourceComment],
     /// The expected number of empty lines before the trailing comments.
-    expected: u32,
+    empty_lines: u32,
 }
 
 impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_> {
@@ -508,7 +520,7 @@ impl Format<PyFormatContext<'_>> for FormatEmptyLinesBeforeTrailingComments<'_>
             .find(|comment| comment.line_position().is_own_line())
         {
             let actual = lines_before(comment.start(), f.context().source()).saturating_sub(1);
-            for _ in actual..self.expected {
+            for _ in actual..self.empty_lines {
                 write!(f, [empty_line()])?;
             }
         }

crates/ruff_python_formatter/src/module/mod_module.rs

@@ -1,23 +1,36 @@
 use ruff_formatter::prelude::hard_line_break;
-use ruff_formatter::write;
+use ruff_formatter::{Buffer, FormatResult};
 use ruff_python_ast::ModModule;
 
-use crate::prelude::*;
+use crate::comments::{trailing_comments, SourceComment};
 use crate::statement::suite::SuiteKind;
+use crate::{write, AsFormat, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
 pub struct FormatModModule;
 
 impl FormatNodeRule<ModModule> for FormatModModule {
     fn fmt_fields(&self, item: &ModModule, f: &mut PyFormatter) -> FormatResult<()> {
         let ModModule { range: _, body } = item;
+        let comments = f.context().comments().clone();
+
         write!(
             f,
             [
                 body.format().with_options(SuiteKind::TopLevel),
+                trailing_comments(comments.dangling(item)),
                 // Trailing newline at the end of the file
                 hard_line_break()
             ]
         )
     }
+
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // Handled as part of `fmt_fields`
+        Ok(())
+    }
 }

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -5,7 +5,6 @@ use ruff_text_size::Ranged;
 
 use crate::comments::format::empty_lines_before_trailing_comments;
 use crate::comments::{leading_comments, trailing_comments, SourceComment};
-use crate::context::NodeLevel;
 use crate::prelude::*;
 use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
 use crate::statement::suite::SuiteKind;
@@ -120,23 +119,15 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
         // # comment
         // ```
         //
-        // At the top-level, reformat as:
+        // At the top-level in a non-stub file, reformat as:
         // ```python
         // class Class:
         //     ...
         //
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(
-            comments.trailing(item),
-            if f.context().node_level() == NodeLevel::TopLevel {
-                2
-            } else {
-                1
-            },
-        )
-        .fmt(f)
+        empty_lines_before_trailing_comments(f, comments.trailing(item)).fmt(f)
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -5,7 +5,6 @@ use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::Ranged;
 
 use crate::comments::SourceComment;
-use crate::context::NodeLevel;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::{Parentheses, Parenthesize};
 use crate::prelude::*;
@@ -156,23 +155,15 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
         // # comment
         // ```
         //
-        // At the top-level, reformat as:
+        // At the top-level in a non-stub file, reformat as:
         // ```python
         // def func():
         //     ...
         //
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(
-            comments.trailing(item),
-            if f.context().node_level() == NodeLevel::TopLevel {
-                2
-            } else {
-                1
-            },
-        )
-        .fmt(f)
+        empty_lines_before_trailing_comments(f, comments.trailing(item)).fmt(f)
     }
 
     fn fmt_dangling_comments(