From af95cbaeef2d654456b19a4e8d60666d786d3dba Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Sat, 28 Oct 2023 03:16:50 +0200
Subject: [PATCH] Add newline after module docstrings in preview style (#8283)

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                    | 39 +++++++++++++------
 .../tests/snapshots/format@preview.py.snap    |  1 +
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs
index c31de9f0f071f..be22098eee650 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,24 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
             )
         } else {
             first.fmt(f)?;
-            (
-                first.statement(),
-                matches!(first, SuiteChildStatement::Docstring(_))
-                    && matches!(self.kind, SuiteKind::Class),
-            )
+
+            #[allow(clippy::if_same_then_else)]
+            let empty_line_after_docstring = if matches!(first, SuiteChildStatement::Docstring(_))
+                && self.kind == SuiteKind::Class
+            {
+                true
+            } else if f.options().preview().is_enabled()
+                && self.kind == SuiteKind::TopLevel
+                && DocstringStmt::try_from_statement(first.statement()).is_some()
+            {
+                // Only in preview mode, insert a newline after a module level docstring, but treat
+                // it as a docstring otherwise. See: https://github.com/psf/black/pull/3932.
+                true
+            } else {
+                false
+            };
+
+            (first.statement(), empty_line_after_docstring)
         };
 
         let mut preceding_comments = comments.leading_dangling_trailing(preceding);
@@ -303,7 +316,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 +402,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                 preceding_comments = following_comments;
             }
 
-            after_class_docstring = false;
+            empty_line_after_docstring = false;
         }
 
         Ok(())
@@ -547,9 +560,11 @@ impl<'a> DocstringStmt<'a> {
         };
 
         if let Expr::Constant(ExprConstant { value, .. }) = value.as_ref() {
-            if !value.is_implicit_concatenated() {
-                return Some(DocstringStmt(stmt));
-            }
+            return match value {
+                Constant::Str(value) if !value.implicit_concatenated => Some(DocstringStmt(stmt)),
+                Constant::Bytes(value) if !value.implicit_concatenated => Some(DocstringStmt(stmt)),
+                _ => None,
+            };
         }
 
         None
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