Standardise newlines after module-level docstrings

CHANGES.md

@@ -22,6 +22,7 @@
 
 - Fix a crash when formatting some dicts with parenthesis-wrapped long string keys
   (#3262)
+- Require one empty line after module-level docstrings (#3287)
 
 ### Configuration
 

docs/the_black_code_style/future_style.md

@@ -131,3 +131,38 @@ with open("bla.txt") as f, open("x"):
 async def main():
     await asyncio.sleep(1)
 ```
+
+### Enforced newline after module docstrings
+
+A single blank line after module docstrings will be enforced, this applies to single and
+multi-line docstrings.
+
+```python
+"""Utility functions and constants."""
+import functools
+```
+
+```python
+"""
+Utility functions and constants.
+"""
+
+
+import functools
+```
+
+will be changed to:
+
+```python
+"""Utility functions and constants."""
+
+import functools
+```
+
+```python
+"""
+Utility functions and constants.
+"""
+
+import functools
+```

src/black/__init__.py

@@ -1073,15 +1073,13 @@ def _format_str_once(src_contents: str, *, mode: Mode) -> str:
     lines = LineGenerator(mode=mode)
     elt = EmptyLineTracker(is_pyi=mode.is_pyi)
     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/linegen.py

@@ -1,6 +1,7 @@
 """
 Generating lines of code.
 """
+
 import sys
 from functools import partial, wraps
 from typing import Collection, Iterator, List, Optional, Set, Union, cast

src/black/lines.py

@@ -1,8 +1,10 @@
 import itertools
 import sys
+from collections import deque
 from dataclasses import dataclass, field
 from typing import (
     Callable,
+    Deque,
     Dict,
     Iterator,
     List,
@@ -451,37 +453,50 @@ 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_defs: List[int] = field(default_factory=list)
-
-    def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
+    # Since we shouldn't need to look back
+    # more than a couple lines we limit
+    # this window to the last 3 visited lines
+    # to keep memory constant when formatting large files.
+    previous_lines_window: Deque[Line] = field(default_factory=deque)
+    previous_lines_window_size = 3
+
+    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_window
+            else before
         )
-        self.previous_after = after
-        self.previous_line = current_line
-        return before, after
+        if (
+            Preview.module_docstring_newlines in current_line.mode
+            and len(self.previous_lines_window) == 1
+            and self.previous_lines_window[-1].is_triple_quoted_string
+        ):
+            before = 1
+
+        self.previous_lines_window.append(current_line)
+        if len(self.previous_lines_window) > self.previous_lines_window_size:
+            self.previous_lines_window.popleft()
+
+        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
@@ -496,8 +511,12 @@ 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_window
+                if (
+                    depth
+                    and not current_line.is_def
+                    and self.previous_lines_window[-1].is_def
+                ):
                     # Empty lines between attributes and methods should be preserved.
                     before = min(1, before)
                 elif depth:
@@ -531,70 +550,80 @@ 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_window
+            and self.previous_lines_window[-1].is_import
             and not current_line.is_import
-            and depth == self.previous_line.depth
+            and depth == self.previous_lines_window[-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_window) > 1
+            and self.previous_lines_window[-2].is_class
+            and self.previous_lines_window[-1].is_triple_quoted_string
+            and current_line.depth == self.previous_lines_window[-1].depth
         ):
-            return before, 1
+            return 1
 
         if (
             Preview.remove_block_trailing_newline in current_line.mode
-            and self.previous_line
-            and self.previous_line.opens_block
+            and self.previous_lines_window
+            and self.previous_lines_window[-1].opens_block
         ):
-            return 0, 0
-        return before, 0
+            return 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_window:
             # 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_window[-1].is_decorator:
+            return 0
 
-            return 0, 0
+        if (
+            self.is_pyi
+            and len(self.previous_lines_window) > 1
+            and self.previous_lines_window[-1].is_stub_class
+            and self.previous_lines_window[-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_window[-1].depth < current_line.depth and (
+            self.previous_lines_window[-1].is_class
+            or self.previous_lines_window[-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_window[-1].is_comment
+            and self.previous_lines_window[-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_window[-1].is_class:
+                if self.previous_lines_window[-1].depth < current_line.depth:
                     newlines = 0
-                elif self.previous_line.depth > current_line.depth:
+                elif self.previous_lines_window[-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_window[-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_window[-1].is_def:
                 if current_line.depth:
                     # In classes empty lines between attributes and methods should
                     # be preserved.
@@ -603,13 +632,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_window[-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

@@ -152,6 +152,7 @@ class Preview(Enum):
     annotation_parens = auto()
     long_docstring_quotes_on_newline = auto()
     normalize_docstring_quotes_and_prefixes_properly = auto()
+    module_docstring_newlines = auto()
     one_element_subscript = auto()
     remove_block_trailing_newline = auto()
     remove_redundant_parens = auto()

src/black/numerics.py

@@ -1,6 +1,7 @@
 """
 Formatting numeric literals.
 """
+
 from blib2to3.pytree import Leaf
 
 

src/black/parsing.py

@@ -1,6 +1,7 @@
 """
 Parse Python code and perform AST validation.
 """
+
 import ast
 import platform
 import sys

src/black/report.py

@@ -1,6 +1,7 @@
 """
 Summarize Black runs to users.
 """
+
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path

src/black/rusty.py

@@ -2,6 +2,7 @@
 
 See https://doc.rust-lang.org/book/ch09-00-error-handling.html.
 """
+
 from typing import Generic, TypeVar, Union
 
 T = TypeVar("T")

src/black/trans.py

@@ -1,6 +1,7 @@
 """
 String transformers that can split and merge strings.
 """
+
 import re
 import sys
 from abc import ABC, abstractmethod

tests/data/miscellaneous/string_quotes.py

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

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