Standardise newlines after module-level docstrings

CHANGES.md

@@ -15,6 +15,7 @@
 <!-- Changes that affect Black's preview style -->
 
 - Remove unnecessary parentheses from `with` statements (#2926)
+- Standardise newlines after module-level docstrings (#2996)
 
 ### _Blackd_
 

fuzz.py

@@ -4,7 +4,6 @@
 generation.  You can run this file with `python`, `pytest`, or (soon)
 a coverage-guided fuzzer I'm working on.
 """
-
 import re
 
 import hypothesmith

scripts/check_pre_commit_rev_in_example.py

@@ -8,7 +8,6 @@
 technical and some pragmatic). Encouraging bad practice is also just
 not ideal. xref: https://github.com/psf/black/issues/420
 """
-
 import os
 import sys
 

scripts/check_version_in_basics_example.py

@@ -3,7 +3,6 @@
 the latest version of Black. This saves us from forgetting to update that
 during the release process.
 """
-
 import os
 import sys
 

scripts/diff_shades_gha_helper.py

@@ -13,7 +13,6 @@
 
 https://black.readthedocs.io/en/latest/contributing/gauging_changes.html#diff-shades
 """
-
 import json
 import os
 import platform

src/black/__init__.py

@@ -1171,17 +1171,15 @@ def _format_str_once(src_contents: str, *, mode: Mode) -> str:
 
     normalize_fmt_off(src_node, preview=mode.preview)
     lines = LineGenerator(mode=mode)
-    elt = EmptyLineTracker(is_pyi=mode.is_pyi)
+    elt = EmptyLineTracker(is_pyi=mode.is_pyi, preview=mode.preview)
     empty_line = Line(mode=mode)
-    after = 0
     split_line_features = {
         feature
         for feature in {Feature.TRAILING_COMMA_IN_CALL, Feature.TRAILING_COMMA_IN_DEF}
         if supports_feature(versions, feature)
     }
     for current_line in lines.visit(src_node):
-        dst_contents.append(str(empty_line) * after)
-        before, after = elt.maybe_empty_lines(current_line)
+        before = elt.maybe_empty_lines(current_line)
         dst_contents.append(str(empty_line) * before)
         for line in transform_line(
             current_line, mode=mode, features=split_line_features

src/black/lines.py

@@ -419,37 +419,47 @@ def __bool__(self) -> bool:
 @dataclass
 class EmptyLineTracker:
     """Provides a stateful method that returns the number of potential extra
-    empty lines needed before and after the currently processed line.
+    empty lines needed before the currently processed line.
 
     Note: this tracker works on lines that haven't been split yet.  It assumes
     the prefix of the first leaf consists of optional newlines.  Those newlines
     are consumed by `maybe_empty_lines()` and included in the computation.
     """
 
     is_pyi: bool = False
-    previous_line: Optional[Line] = None
-    previous_after: int = 0
+    previous_lines: List[Line] = field(default_factory=list)
     previous_defs: List[int] = field(default_factory=list)
+    preview: bool = False
 
-    def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
+    def maybe_empty_lines(self, current_line: Line) -> int:
         """Return the number of extra empty lines before and after the `current_line`.
 
         This is for separating `def`, `async def` and `class` with extra empty
         lines (two on module-level).
         """
-        before, after = self._maybe_empty_lines(current_line)
+        before = self._maybe_empty_lines(current_line)
         before = (
             # Black should not insert empty lines at the beginning
             # of the file
             0
-            if self.previous_line is None
-            else before - self.previous_after
+            if not self.previous_lines
+            else before
         )
-        self.previous_after = after
-        self.previous_line = current_line
-        return before, after
+        if (
+            self.preview
+            and len(self.previous_lines) == 1
+            and self.previous_lines[-1].is_triple_quoted_string
+        ):
+            # Newlines after multi-line module level docstring.
+            if str(self.previous_lines[-1].leaves[0]).count("\n") >= 1:
+                before = 0
+            else:
+                before = 1
+
+        self.previous_lines.append(current_line)
+        return before
 
-    def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
+    def _maybe_empty_lines(self, current_line: Line) -> int:
         max_allowed = 1
         if current_line.depth == 0:
             max_allowed = 1 if self.is_pyi else 2
@@ -464,8 +474,8 @@ def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
         depth = current_line.depth
         while self.previous_defs and self.previous_defs[-1] >= depth:
             if self.is_pyi:
-                assert self.previous_line is not None
-                if depth and not current_line.is_def and self.previous_line.is_def:
+                assert self.previous_lines
+                if depth and not current_line.is_def and self.previous_lines[-1].is_def:
                     # Empty lines between attributes and methods should be preserved.
                     before = min(1, before)
                 elif depth:
@@ -499,64 +509,72 @@ def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
             return self._maybe_empty_lines_for_class_or_def(current_line, before)
 
         if (
-            self.previous_line
-            and self.previous_line.is_import
+            self.previous_lines
+            and self.previous_lines[-1].is_import
             and not current_line.is_import
-            and depth == self.previous_line.depth
+            and depth == self.previous_lines[-1].depth
         ):
-            return (before or 1), 0
+            return before or 1
 
         if (
-            self.previous_line
-            and self.previous_line.is_class
-            and current_line.is_triple_quoted_string
+            len(self.previous_lines) > 1
+            and self.previous_lines[-2].is_class
+            and self.previous_lines[-1].is_triple_quoted_string
+            and current_line.depth == self.previous_lines[-1].depth
         ):
-            return before, 1
+            return 1
 
-        return before, 0
+        return before
 
     def _maybe_empty_lines_for_class_or_def(
         self, current_line: Line, before: int
-    ) -> Tuple[int, int]:
+    ) -> int:
         if not current_line.is_decorator:
             self.previous_defs.append(current_line.depth)
-        if self.previous_line is None:
+        if not self.previous_lines:
             # Don't insert empty lines before the first line in the file.
-            return 0, 0
+            return 0
 
-        if self.previous_line.is_decorator:
-            if self.is_pyi and current_line.is_stub_class:
-                # Insert an empty line after a decorated stub class
-                return 0, 1
+        if self.previous_lines[-1].is_decorator:
+            return 0
 
-            return 0, 0
+        if (
+            self.is_pyi
+            and len(self.previous_lines) > 1
+            and self.previous_lines[-1].is_stub_class
+            and self.previous_lines[-2].is_decorator
+        ):
+            # Insert an empty line after a decorated stub class
+            return 1
 
-        if self.previous_line.depth < current_line.depth and (
-            self.previous_line.is_class or self.previous_line.is_def
+        if self.previous_lines[-1].depth < current_line.depth and (
+            self.previous_lines[-1].is_class or self.previous_lines[-1].is_def
         ):
-            return 0, 0
+            return 0
 
         if (
-            self.previous_line.is_comment
-            and self.previous_line.depth == current_line.depth
+            self.previous_lines[-1].is_comment
+            and self.previous_lines[-1].depth == current_line.depth
             and before == 0
         ):
-            return 0, 0
+            return 0
 
         if self.is_pyi:
-            if current_line.is_class or self.previous_line.is_class:
-                if self.previous_line.depth < current_line.depth:
+            if current_line.is_class or self.previous_lines[-1].is_class:
+                if self.previous_lines[-1].depth < current_line.depth:
                     newlines = 0
-                elif self.previous_line.depth > current_line.depth:
+                elif self.previous_lines[-1].depth > current_line.depth:
                     newlines = 1
-                elif current_line.is_stub_class and self.previous_line.is_stub_class:
+                elif (
+                    current_line.is_stub_class and self.previous_lines[-1].is_stub_class
+                ):
                     # No blank line between classes with an empty body
                     newlines = 0
                 else:
                     newlines = 1
             elif (
                 current_line.is_def or current_line.is_decorator
-            ) and not self.previous_line.is_def:
+            ) and not self.previous_lines[-1].is_def:
                 if current_line.depth:
                     # In classes empty lines between attributes and methods should
                     # be preserved.
@@ -565,13 +583,13 @@ def _maybe_empty_lines_for_class_or_def(
                     # Blank line between a block of functions (maybe with preceding
                     # decorators) and a block of non-functions
                     newlines = 1
-            elif self.previous_line.depth > current_line.depth:
+            elif self.previous_lines[-1].depth > current_line.depth:
                 newlines = 1
             else:
                 newlines = 0
         else:
             newlines = 1 if current_line.depth else 2
-        return newlines, 0
+        return newlines
 
 
 def enumerate_reversed(sequence: Sequence[T]) -> Iterator[Tuple[Index, T]]:

src/black/mode.py

@@ -3,7 +3,6 @@
 Mostly around Python language feature support per version and Black configuration
 chosen by the user.
 """
-
 from hashlib import sha256
 import sys
 

src/black/nodes.py

@@ -1,7 +1,6 @@
 """
 blib2to3 Node/Leaf transformation-related utility functions.
 """
-
 import sys
 from typing import (
     Generic,

src/black/output.py

@@ -2,7 +2,6 @@
 
 The double calls are for patching purposes in tests.
 """
-
 import json
 from typing import Any, Optional
 from mypy_extensions import mypyc_attr

src/black/strings.py

@@ -1,7 +1,6 @@
 """
 Simple formatting on strings. Further string formatting code is in trans.py.
 """
-
 import re
 import sys
 from functools import lru_cache

tests/data/module_docstring_1.py

@@ -0,0 +1,25 @@
+"""Single line module-level docstring should be followed by single newline."""
+
+
+
+
+a = 1
+
+
+"""I'm just a string so should be followed by 2 newlines."""
+
+
+
+
+b = 2
+
+# output
+"""Single line module-level docstring should be followed by single newline."""
+
+a = 1
+
+
+"""I'm just a string so should be followed by 2 newlines."""
+
+
+b = 2

tests/data/module_docstring_2.py

@@ -0,0 +1,66 @@
+"""I am a very helpful module docstring.
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit,
+sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+Ut enim ad minim veniam,
+quis nostrud exercitation ullamco laboris
+nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate
+velit esse cillum dolore eu fugiat nulla pariatur.
+Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+
+
+
+
+a = 1
+
+
+"""Look at me I'm a docstring...
+
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+........................................................NOT!
+"""
+
+
+
+
+b = 2
+
+# output
+"""I am a very helpful module docstring.
+
+Lorem ipsum dolor sit amet, consectetur adipiscing elit,
+sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+Ut enim ad minim veniam,
+quis nostrud exercitation ullamco laboris
+nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate
+velit esse cillum dolore eu fugiat nulla pariatur.
+Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+a = 1
+
+
+"""Look at me I'm a docstring...
+
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+............................................................
+........................................................NOT!
+"""
+
+
+b = 2

tests/data/string_quotes.py

@@ -1,4 +1,5 @@
 ''''''
+
 '\''
 '"'
 "'"
@@ -57,8 +58,8 @@
 f"\"{a}\"{'hello' * b}\"{c}\""
 
 # output
-
 """"""
+
 "'"
 '"'
 "'"

tests/optional.py

@@ -13,7 +13,6 @@
 
 Adapted from https://pypi.org/project/pytest-optional-tests/, (c) 2019 Reece Hart
 """
-
 from functools import lru_cache
 import itertools
 import logging