No newline after function docstrings (#8375)

konstin · web-flow · commit 3076d76b0ab7 · 2023-10-31T14:32:15.000-04:00
Fixup for #8216 to not apply to function docstrings. Main before #8216: | project | similarity index | total files | changed files | |----------------|------------------:|------------------:|------------------:| | cpython | 0.75804 | 1799 | 1648 | | django | 0.99984 | 2772 | 33 | | home-assistant | 0.99963 | 10596 | 148 | | poetry | 0.99925 | 317 | 12 | | transformers | 0.99967 | 2657 | 328 | | twine | 1.00000 | 33 | 0 | | typeshed | 0.99980 | 3669 | 18 | | warehouse | 0.99977 | 654 | 13 | | zulip | 0.99970 | 1459 | 22 | main now: | project | similarity index | total files | changed files | |----------------|------------------:|------------------:|------------------:| | cpython | 0.75804 | 1799 | 1648 | | django | 0.99984 | 2772 | 48 | | home-assistant | 0.99963 | 10596 | 181 | | poetry | 0.99925 | 317 | 12 | | transformers | 0.99967 | 2657 | 339 | | twine | 1.00000 | 33 | 0 | | typeshed | 0.99980 | 3669 | 18 | | warehouse | 0.99977 | 654 | 13 | | zulip | 0.99970 | 1459 | 23 | PR: | project | similarity index | total files | changed files | |----------------|------------------:|------------------:|------------------:| | cpython | 0.75804 | 1799 | 1648 | | django | 0.99984 | 2772 | 33 | | home-assistant | 0.99963 | 10596 | 148 | | poetry | 0.99925 | 317 | 12 | | transformers | 0.99967 | 2657 | 328 | | twine | 1.00000 | 33 | 0 | | typeshed | 0.99980 | 3669 | 18 | | warehouse | 0.99977 | 654 | 13 | | zulip | 0.99970 | 1459 | 22 |
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/docstring.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/docstring.py
@@ -136,6 +136,13 @@ class CommentAfterDocstring5:
     # This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+    """Browse module classes and functions in IDLE."""
+    # ^ Do not insert a newline above here
+
+    pass
+
+
 class TabbedIndent:
 	def tabbed_indent(self):
 		"""check for correct tabbed formatting
diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs
@@ -89,15 +89,15 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
             }
 
             SuiteKind::Function => {
-                if let Some(docstring) = DocstringStmt::try_from_statement(first) {
+                if let Some(docstring) = DocstringStmt::try_from_statement(first, self.kind) {
                     SuiteChildStatement::Docstring(docstring)
                 } else {
                     SuiteChildStatement::Other(first)
                 }
             }
 
             SuiteKind::Class => {
-                if let Some(docstring) = DocstringStmt::try_from_statement(first) {
+                if let Some(docstring) = DocstringStmt::try_from_statement(first, self.kind) {
                     if !comments.has_leading(first)
                         && lines_before(first.start(), source) > 1
                         && !source_type.is_stub()
@@ -150,7 +150,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                 true
             } else if f.options().preview().is_enabled()
                 && self.kind == SuiteKind::TopLevel
-                && DocstringStmt::try_from_statement(first.statement()).is_some()
+                && DocstringStmt::try_from_statement(first.statement(), self.kind).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.
@@ -543,17 +543,25 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for Suite {
 
 /// A statement representing a docstring.
 #[derive(Copy, Clone, Debug)]
-pub(crate) struct DocstringStmt<'a>(&'a Stmt);
+pub(crate) struct DocstringStmt<'a> {
+    /// The [`Stmt::Expr`]
+    docstring: &'a Stmt,
+    /// The parent suite kind
+    suite_kind: SuiteKind,
+}
 
 impl<'a> DocstringStmt<'a> {
     /// Checks if the statement is a simple string that can be formatted as a docstring
-    fn try_from_statement(stmt: &'a Stmt) -> Option<DocstringStmt<'a>> {
+    fn try_from_statement(stmt: &'a Stmt, suite_kind: SuiteKind) -> Option<DocstringStmt<'a>> {
         let Stmt::Expr(ast::StmtExpr { value, .. }) = stmt else {
             return None;
         };
 
         match value.as_ref() {
-            Expr::StringLiteral(value) if !value.implicit_concatenated => Some(DocstringStmt(stmt)),
+            Expr::StringLiteral(value) if !value.implicit_concatenated => Some(DocstringStmt {
+                docstring: stmt,
+                suite_kind,
+            }),
             _ => None,
         }
     }
@@ -562,14 +570,14 @@ impl<'a> DocstringStmt<'a> {
 impl Format<PyFormatContext<'_>> for DocstringStmt<'_> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
         let comments = f.context().comments().clone();
-        let node_comments = comments.leading_dangling_trailing(self.0);
+        let node_comments = comments.leading_dangling_trailing(self.docstring);
 
         if FormatStmtExpr.is_suppressed(node_comments.trailing, f.context()) {
-            suppressed_node(self.0).fmt(f)
+            suppressed_node(self.docstring).fmt(f)
         } else {
             // SAFETY: Safe because `DocStringStmt` guarantees that it only ever wraps a `ExprStmt` containing a `ExprStringLiteral`.
             let string_literal = self
-                .0
+                .docstring
                 .as_expr_stmt()
                 .unwrap()
                 .value
@@ -587,23 +595,25 @@ impl Format<PyFormatContext<'_>> for DocstringStmt<'_> {
                 ]
             )?;
 
-            // Comments after docstrings need a newline between the docstring and the comment.
-            // (https://github.com/astral-sh/ruff/issues/7948)
-            // ```python
-            // class ModuleBrowser:
-            //     """Browse module classes and functions in IDLE."""
-            //     # ^ Insert a newline above here
-            //
-            //     def __init__(self, master, path, *, _htest=False, _utest=False):
-            //         pass
-            // ```
-            if let Some(own_line) = node_comments
-                .trailing
-                .iter()
-                .find(|comment| comment.line_position().is_own_line())
-            {
-                if lines_before(own_line.start(), f.context().source()) < 2 {
-                    empty_line().fmt(f)?;
+            if self.suite_kind == SuiteKind::Class {
+                // Comments after class docstrings need a newline between the docstring and the
+                // comment (https://github.com/astral-sh/ruff/issues/7948).
+                // ```python
+                // class ModuleBrowser:
+                //     """Browse module classes and functions in IDLE."""
+                //     # ^ Insert a newline above here
+                //
+                //     def __init__(self, master, path, *, _htest=False, _utest=False):
+                //         pass
+                // ```
+                if let Some(own_line) = node_comments
+                    .trailing
+                    .iter()
+                    .find(|comment| comment.line_position().is_own_line())
+                {
+                    if lines_before(own_line.start(), f.context().source()) < 2 {
+                        empty_line().fmt(f)?;
+                    }
                 }
             }
 
@@ -625,7 +635,7 @@ pub(crate) enum SuiteChildStatement<'a> {
 impl<'a> SuiteChildStatement<'a> {
     pub(crate) const fn statement(self) -> &'a Stmt {
         match self {
-            SuiteChildStatement::Docstring(docstring) => docstring.0,
+            SuiteChildStatement::Docstring(docstring) => docstring.docstring,
             SuiteChildStatement::Other(statement) => statement,
         }
     }
diff --git a/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__fmtonoff.py.snap b/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__fmtonoff.py.snap
@@ -225,11 +225,8 @@ d={'a':1,
          # fmt: on
          goes + here,
          andhere,
-@@ -120,10 +121,13 @@
- 
-     The comments between will be formatted. This is a known limitation.
+@@ -122,8 +123,10 @@
      """
-+
      # fmt: off
  
 -    # hey, that won't work
@@ -240,7 +237,7 @@ d={'a':1,
      # fmt: on
      pass
  
-@@ -138,7 +142,7 @@
+@@ -138,7 +141,7 @@
      now . considers . multiple . fmt . directives . within . one . prefix
      # fmt: on
      # fmt: off
@@ -249,7 +246,7 @@ d={'a':1,
      # fmt: on
  
  
-@@ -178,14 +182,18 @@
+@@ -178,14 +181,18 @@
          $
          """,
          # fmt: off
@@ -398,7 +395,6 @@ def off_and_on_without_data():
 
     The comments between will be formatted. This is a known limitation.
     """
-
     # fmt: off
 
 
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@docstring.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@docstring.py.snap
@@ -142,6 +142,13 @@ class CommentAfterDocstring5:
     # This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+    """Browse module classes and functions in IDLE."""
+    # ^ Do not insert a newline above here
+
+    pass
+
+
 class TabbedIndent:
 	def tabbed_indent(self):
 		"""check for correct tabbed formatting
@@ -301,6 +308,13 @@ class CommentAfterDocstring5:
     # This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+    """Browse module classes and functions in IDLE."""
+    # ^ Do not insert a newline above here
+
+    pass
+
+
 class TabbedIndent:
     def tabbed_indent(self):
         """check for correct tabbed formatting
@@ -460,6 +474,13 @@ class CommentAfterDocstring5:
   # This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+  """Browse module classes and functions in IDLE."""
+  # ^ Do not insert a newline above here
+
+  pass
+
+
 class TabbedIndent:
   def tabbed_indent(self):
     """check for correct tabbed formatting
@@ -619,6 +640,13 @@ class CommentAfterDocstring5:
 	# This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+	"""Browse module classes and functions in IDLE."""
+	# ^ Do not insert a newline above here
+
+	pass
+
+
 class TabbedIndent:
 	def tabbed_indent(self):
 		"""check for correct tabbed formatting
@@ -778,6 +806,13 @@ class CommentAfterDocstring5:
 	# This class is also the base class for pathbrowser.PathBrowser.
 
 
+def f():
+	"""Browse module classes and functions in IDLE."""
+	# ^ Do not insert a newline above here
+
+	pass
+
+
 class TabbedIndent:
 	def tabbed_indent(self):
 		"""check for correct tabbed formatting
