Keep trailing new line in module docstring

astral-sh · Feb 1, 2024 · 7a13f96 · 7a13f96
1 parent 0e8d1b2
commit 7a13f96
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 156 deletions.
diff --git a/crates/ruff_python_formatter/src/other/string_literal.rs b/crates/ruff_python_formatter/src/other/string_literal.rs
@@ -66,7 +66,11 @@ impl Format<PyFormatContext<'_>> for FormatStringLiteral<'_> {
         );
 
         if self.layout.is_docstring() {
-            docstring::format(&normalized, f)
+            let is_module = matches!(
+                f.context().node_level(),
+                crate::context::NodeLevel::TopLevel(_)
+            );
+            docstring::format(&normalized, f, is_module)
         } else {
             normalized.fmt(f)
         }

diff --git a/crates/ruff_python_formatter/src/string/docstring.rs b/crates/ruff_python_formatter/src/string/docstring.rs
@@ -102,7 +102,11 @@ use super::{NormalizedString, QuoteChar};
 ///         line c
 ///    """
 /// ```
-pub(crate) fn format(normalized: &NormalizedString, f: &mut PyFormatter) -> FormatResult<()> {
+pub(crate) fn format(
+    normalized: &NormalizedString,
+    f: &mut PyFormatter,
+    is_module: bool,
+) -> FormatResult<()> {
     let docstring = &normalized.text;
 
     // Black doesn't change the indentation of docstrings that contain an escaped newline
@@ -205,6 +209,10 @@ pub(crate) fn format(normalized: &NormalizedString, f: &mut PyFormatter) -> Form
         space().fmt(f)?;
     }
 
+    if is_module && trim_end.ends_with('\n') {
+        hard_line_break().fmt(f)?;
+    }
+
     write!(f, [source_position(normalized.end()), normalized.quotes])
 }
 

diff --git a/...ff_python_formatter/tests/snapshots/black_compatibility@cases__module_docstring_2.py.snap b/...ff_python_formatter/tests/snapshots/black_compatibility@cases__module_docstring_2.py.snap
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@preview.py.snap
@@ -189,13 +189,10 @@ def f():
 ```diff
 --- Stable
 +++ Preview
-@@ -13,15 +13,14 @@
-     def callback(event         ):
-         print('Python:Click')
+@@ -15,13 +15,13 @@
  
--    button.on_event(ButtonClick, callback)
--"""
-+    button.on_event(ButtonClick, callback)"""
+     button.on_event(ButtonClick, callback)
+ """
 +
  first_stmt_after_module_level_docstring = 1
  
@@ -208,7 +205,7 @@ def f():
  
  
  def raw_docstring():
-@@ -41,23 +40,22 @@
+@@ -41,23 +41,22 @@
  
  
  class RemoveNewlineBeforeClassDocstring:
@@ -241,7 +238,7 @@ def f():
  
      aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb] = (
          cccccccc.ccccccccccccc(d).cccccccc + e
-@@ -71,9 +69,9 @@
+@@ -71,9 +70,9 @@
          + eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee
      )
  
@@ -291,7 +288,8 @@ A code black to format
     def callback(event         ):
         print('Python:Click')
 
-    button.on_event(ButtonClick, callback)"""
+    button.on_event(ButtonClick, callback)
+"""
 
 first_stmt_after_module_level_docstring = 1
 

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@stub_files__suite.pyi.snap b/crates/ruff_python_formatter/tests/snapshots/format@stub_files__suite.pyi.snap
@@ -313,24 +313,15 @@ class ComplexStatements:
 ```diff
 --- Stable
 +++ Preview
-@@ -1,7 +1,6 @@
- """Tests for empty line rules in stub files, mostly inspired by typeshed.
- The rules are a list of nested exceptions. See also
--https://github.com/psf/black/blob/c160e4b7ce30c661ac4f2dfa5038becf1b8c5c33/src/black/lines.py#L576-L744
--"""
-+https://github.com/psf/black/blob/c160e4b7ce30c661ac4f2dfa5038becf1b8c5c33/src/black/lines.py#L576-L744"""
- 
- import sys
- from typing import Self, TypeAlias, final
-@@ -110,6 +109,7 @@
+@@ -110,6 +110,7 @@
  
      class InnerClass5:
          def a(self): ...
 +
      field1 = 1
  
      class InnerClass6:
-@@ -119,6 +119,7 @@
+@@ -119,6 +120,7 @@
  
      class InnerClass7:
          def a(self): ...