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 #7995

Author: konstin

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

crates/ruff_python_formatter/tests/snapshots/[email protected]

@@ -166,6 +166,7 @@ preview                 = Enabled
 """
 Black's `Preview.module_docstring_newlines`
 """
+
 first_stmt_after_module_level_docstring = 1