Format empty lines in stub files closer to 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/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

@@ -1,5 +1,5 @@
 use ruff_formatter::write;
-use ruff_python_ast::{Decorator, StmtClassDef};
+use ruff_python_ast::{Decorator, PySourceType, StmtClassDef};
 use ruff_python_trivia::lines_after_ignoring_trivia;
 use ruff_text_size::Ranged;
 
@@ -128,15 +128,13 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(
-            comments.trailing(item),
-            if f.context().node_level() == NodeLevel::TopLevel {
-                2
-            } else {
-                1
-            },
-        )
-        .fmt(f)
+        let empty_lines = match (f.options().source_type(), f.context().node_level()) {
+            (PySourceType::Stub, NodeLevel::TopLevel) => 1,
+            (PySourceType::Stub, _) => 0,
+            (_, NodeLevel::TopLevel) => 2,
+            (_, _) => 1,
+        };
+        empty_lines_before_trailing_comments(comments.trailing(item), empty_lines).fmt(f)
     }
 
     fn fmt_dangling_comments(

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -1,6 +1,6 @@
 use crate::comments::format::empty_lines_before_trailing_comments;
 use ruff_formatter::write;
-use ruff_python_ast::{Parameters, StmtFunctionDef};
+use ruff_python_ast::{Parameters, PySourceType, StmtFunctionDef};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::Ranged;
 
@@ -164,15 +164,13 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
         //
         // # comment
         // ```
-        empty_lines_before_trailing_comments(
-            comments.trailing(item),
-            if f.context().node_level() == NodeLevel::TopLevel {
-                2
-            } else {
-                1
-            },
-        )
-        .fmt(f)
+        let empty_lines = match (f.options().source_type(), f.context().node_level()) {
+            (PySourceType::Stub, NodeLevel::TopLevel) => 1,
+            (PySourceType::Stub, _) => 0,
+            (_, NodeLevel::TopLevel) => 2,
+            (_, _) => 1,
+        };
+        empty_lines_before_trailing_comments(comments.trailing(item), empty_lines).fmt(f)
     }
 
     fn fmt_dangling_comments(