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