Unify docstring detection (psf#4095)

Co-authored-by: hauntsaninja <hauntsaninja@gmail.com>

CHANGES.md

@@ -14,6 +14,7 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Format module docstrings the same as class and function docstrings (#4095)
 - Fix bug where spaces were not added around parenthesized walruses in subscripts,
   unlike other binary operators (#4109)
 

src/black/linegen.py

@@ -424,7 +424,7 @@ def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
         if Preview.hex_codes_in_unicode_sequences in self.mode:
             normalize_unicode_escape_sequences(leaf)
 
-        if is_docstring(leaf) and not re.search(r"\\\s*\n", leaf.value):
+        if is_docstring(leaf, self.mode) and not re.search(r"\\\s*\n", leaf.value):
             # We're ignoring docstrings with backslash newline escapes because changing
             # indentation of those changes the AST representation of the code.
             if self.mode.string_normalization:
@@ -477,7 +477,7 @@ def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
             quote = quote_char * quote_len
 
             # It's invalid to put closing single-character quotes on a new line.
-            if self.mode and quote_len == 3:
+            if quote_len == 3:
                 # We need to find the length of the last line of the docstring
                 # to find if we can add the closing quotes to the line without
                 # exceeding the maximum line length.

src/black/lines.py

@@ -196,7 +196,7 @@ def is_class_paren_empty(self) -> bool:
         )
 
     @property
-    def is_triple_quoted_string(self) -> bool:
+    def _is_triple_quoted_string(self) -> bool:
         """Is the line a triple quoted string?"""
         if not self or self.leaves[0].type != token.STRING:
             return False
@@ -209,6 +209,13 @@ def is_triple_quoted_string(self) -> bool:
             return True
         return False
 
+    @property
+    def is_docstring(self) -> bool:
+        """Is the line a docstring?"""
+        if Preview.unify_docstring_detection not in self.mode:
+            return self._is_triple_quoted_string
+        return bool(self) and is_docstring(self.leaves[0], self.mode)
+
     @property
     def is_chained_assignment(self) -> bool:
         """Is the line a chained assignment"""
@@ -583,7 +590,7 @@ def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
             and self.previous_block
             and self.previous_block.previous_block is None
             and len(self.previous_block.original_line.leaves) == 1
-            and self.previous_block.original_line.is_triple_quoted_string
+            and self.previous_block.original_line.is_docstring
             and not (current_line.is_class or current_line.is_def)
         ):
             before = 1
@@ -690,7 +697,7 @@ def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
         if (
             self.previous_line
             and self.previous_line.is_class
-            and current_line.is_triple_quoted_string
+            and current_line.is_docstring
         ):
             if Preview.no_blank_line_before_class_docstring in current_line.mode:
                 return 0, 1
@@ -701,7 +708,7 @@ def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
         is_empty_first_line_ok = (
             Preview.allow_empty_first_line_in_block in current_line.mode
             and (
-                not is_docstring(current_line.leaves[0])
+                not is_docstring(current_line.leaves[0], current_line.mode)
                 or (
                     self.previous_line
                     and self.previous_line.leaves[0]

src/black/mode.py

@@ -195,6 +195,7 @@ class Preview(Enum):
     single_line_format_skip_with_multiple_comments = auto()
     long_case_block_line_splitting = auto()
     allow_form_feeds = auto()
+    unify_docstring_detection = auto()
     respect_east_asian_width = auto()
 
 

src/black/nodes.py

@@ -531,14 +531,24 @@ def is_arith_like(node: LN) -> bool:
     }
 
 
-def is_docstring(leaf: Leaf) -> bool:
+def is_docstring(leaf: Leaf, mode: Mode) -> bool:
     if leaf.type != token.STRING:
         return False
 
     prefix = get_string_prefix(leaf.value)
     if set(prefix).intersection("bBfF"):
         return False
 
+    if (
+        Preview.unify_docstring_detection in mode
+        and leaf.parent
+        and leaf.parent.type == syms.simple_stmt
+        and not leaf.parent.prev_sibling
+        and leaf.parent.parent
+        and leaf.parent.parent.type == syms.file_input
+    ):
+        return True
+
     if prev_siblings_are(
         leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
     ):

src/black/strings.py

@@ -63,6 +63,8 @@ def lines_with_leading_tabs_expanded(s: str) -> List[str]:
             )
         else:
             lines.append(line)
+    if s.endswith("\n"):
+        lines.append("")
     return lines
 
 

tests/data/cases/module_docstring_2.py

@@ -1,6 +1,7 @@
 # flags: --preview
 """I am a very helpful module docstring.
 
+With trailing spaces:      
 Lorem ipsum dolor sit amet, consectetur adipiscing elit,
 sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
 Ut enim ad minim veniam,
@@ -38,6 +39,7 @@
 # output
 """I am a very helpful module docstring.
 
+With trailing spaces:
 Lorem ipsum dolor sit amet, consectetur adipiscing elit,
 sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
 Ut enim ad minim veniam,

tests/data/cases/preview_no_blank_line_before_docstring.py

@@ -29,6 +29,9 @@ class MultilineDocstringsAsWell:
     and on so many lines...
     """
 
+class SingleQuotedDocstring:
+
+    "I'm a docstring but I don't even get triple quotes."
 
 # output
 
@@ -57,3 +60,7 @@ class MultilineDocstringsAsWell:
 
     and on so many lines...
     """
+
+
+class SingleQuotedDocstring:
+    "I'm a docstring but I don't even get triple quotes."