From 43558f551303e804db73fd04b3fe7a1681793093 Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Fri, 27 Oct 2023 14:58:12 +0200
Subject: [PATCH] Add newline after module docstrings in preview style

Change
```python
"""Test docstring"""
a = 1
```
to
```python
"""Test docstring"""

a = 1
```
in preview style, but don't touch the docstring otherwise.

Do we want to ask black to also format the content of module level docstrings? Seems inconsistent to me that we change function and class docstring indentation/contents but not module docstrings.

Fixes https://github.com/astral-sh/ruff/issues/7995
---
 .../src/statement/suite.rs                    | 30 +++++++++++++------
 .../format@expression__number.py.snap         | 13 ++++++++
 .../tests/snapshots/format@preview.py.snap    |  1 +
 3 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs
index c31de9f0f071f..ee95bef96e185 100644
--- a/crates/ruff_python_formatter/src/statement/suite.rs
+++ b/crates/ruff_python_formatter/src/statement/suite.rs
@@ -19,7 +19,7 @@ use crate::verbatim::{
 };
 
 /// Level at which the [`Suite`] appears in the source code.
-#[derive(Copy, Clone, Debug, Default)]
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
 pub enum SuiteKind {
     /// Statements at the module level / top level
     TopLevel,
@@ -123,7 +123,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
 
         let first_comments = comments.leading_dangling_trailing(first);
 
-        let (mut preceding, mut after_class_docstring) = if first_comments
+        let (mut preceding, mut empty_line_after_docstring) = if first_comments
             .leading
             .iter()
             .any(|comment| comment.is_suppression_off_comment(source))
@@ -143,11 +143,23 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
             )
         } else {
             first.fmt(f)?;
-            (
-                first.statement(),
-                matches!(first, SuiteChildStatement::Docstring(_))
-                    && matches!(self.kind, SuiteKind::Class),
-            )
+            let empty_line_after_docstring = if matches!(first, SuiteChildStatement::Docstring(_))
+                && self.kind == SuiteKind::Class
+            {
+                true
+            } else if f.options().preview().is_enabled()
+                && DocstringStmt::try_from_statement(first.statement()).is_some()
+                && self.kind == SuiteKind::TopLevel
+            {
+                // Only in preview mode, insert a newline after a module level docstring, but treat
+                // it as a docstring otherwise
+                // https://github.com/psf/black/pull/3932
+                // https://github.com/astral-sh/ruff/issues/7995
+                true
+            } else {
+                false
+            };
+            (first.statement(), empty_line_after_docstring)
         };
 
         let mut preceding_comments = comments.leading_dangling_trailing(preceding);
@@ -303,7 +315,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                         }
                     },
                 }
-            } else if after_class_docstring {
+            } else if empty_line_after_docstring {
                 // Enforce an empty line after a class docstring, e.g., these are both stable
                 // formatting:
                 // ```python
@@ -389,7 +401,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                 preceding_comments = following_comments;
             }
 
-            after_class_docstring = false;
+            empty_line_after_docstring = false;
         }
 
         Ok(())
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__number.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__number.py.snap
index feea8103e080f..02c79b23a4204 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@expression__number.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__number.py.snap
@@ -25,4 +25,17 @@ input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression
 ```
 
 
+## Preview changes
+```diff
+--- Stable
++++ Preview
+@@ -1,4 +1,5 @@
+ 0.1
++
+ 1.0
+ 1e1
+ 1e-1
+```
+
+
 
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap
index b803c6f3fd74b..f37b61a77bf1c 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap
@@ -166,6 +166,7 @@ preview                 = Enabled
 """
 Black's `Preview.module_docstring_newlines`
 """
+
 first_stmt_after_module_level_docstring = 1