Put closing quote on a separate line if docstring is too long (#3044)

Fixes #1632

Co-authored-by: Felix Hildén <felix.hilden@gmail.com>

CHANGES.md

@@ -17,6 +17,7 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fixed bug where docstrings with triple quotes could exceed max line length (#3044)
 - Remove redundant parentheses around awaited objects (#2991)
 - Parentheses around return annotations are now managed (#2990)
 - Remove unnecessary parentheses from `with` statements (#2926)

src/black/linegen.py

@@ -305,9 +305,9 @@ def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
             quote_len = 1 if docstring[1] != quote_char else 3
             docstring = docstring[quote_len:-quote_len]
             docstring_started_empty = not docstring
+            indent = " " * 4 * self.current_line.depth
 
             if is_multiline_string(leaf):
-                indent = " " * 4 * self.current_line.depth
                 docstring = fix_docstring(docstring, indent)
             else:
                 docstring = docstring.strip()
@@ -329,6 +329,28 @@ def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
 
             # We could enforce triple quotes at this point.
             quote = quote_char * quote_len
+
+            if Preview.long_docstring_quotes_on_newline in self.mode:
+                # 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.
+                # If docstring is one line, then we need to add the length
+                # of the indent, prefix, and starting quotes. Ending quote are
+                # handled later
+                lines = docstring.splitlines()
+                last_line_length = len(lines[-1]) if docstring else 0
+
+                if len(lines) == 1:
+                    last_line_length += len(indent) + len(prefix) + quote_len
+
+                # If adding closing quotes would cause the last line to exceed
+                # the maximum line length then put a line break before the
+                # closing quotes
+                if last_line_length + quote_len > self.mode.line_length:
+                    leaf.value = prefix + quote + docstring + "\n" + indent + quote
+                else:
+                    leaf.value = prefix + quote + docstring + quote
+            else:
                 leaf.value = prefix + quote + docstring + quote
 
         yield from self.visit_default(leaf)

src/black/mode.py

@@ -147,6 +147,7 @@ class Preview(Enum):
     remove_redundant_parens = auto()
     one_element_subscript = auto()
     annotation_parens = auto()
+    long_docstring_quotes_on_newline = auto()
 
 
 class Deprecated(UserWarning):

tests/data/docstring.py

@@ -188,6 +188,27 @@ def my_god_its_full_of_stars_2():
     "I'm sorry Dave "
 
 
+def docstring_almost_at_line_limit():
+    """long docstring................................................................."""
+
+
+def docstring_almost_at_line_limit2():
+    """long docstring.................................................................
+
+    ..................................................................................
+    """
+
+
+def docstring_at_line_limit():
+    """long docstring................................................................"""
+
+
+def multiline_docstring_at_line_limit():
+    """first line-----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""
+
+
 # output
 
 class MyClass:
@@ -375,3 +396,24 @@ def my_god_its_full_of_stars_1():
 # the space below is actually a \u2001, removed in output
 def my_god_its_full_of_stars_2():
     "I'm sorry Dave"
+
+
+def docstring_almost_at_line_limit():
+    """long docstring................................................................."""
+
+
+def docstring_almost_at_line_limit2():
+    """long docstring.................................................................
+
+    ..................................................................................
+    """
+
+
+def docstring_at_line_limit():
+    """long docstring................................................................"""
+
+
+def multiline_docstring_at_line_limit():
+    """first line-----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""

tests/data/docstring_preview.py

@@ -0,0 +1,89 @@
+def docstring_almost_at_line_limit():
+    """long docstring.................................................................
+    """
+
+
+def docstring_almost_at_line_limit_with_prefix():
+    f"""long docstring................................................................
+    """
+
+
+def mulitline_docstring_almost_at_line_limit():
+    """long docstring.................................................................
+
+    ..................................................................................
+    """
+
+
+def mulitline_docstring_almost_at_line_limit_with_prefix():
+    f"""long docstring................................................................
+
+    ..................................................................................
+    """
+
+
+def docstring_at_line_limit():
+    """long docstring................................................................"""
+
+
+def docstring_at_line_limit_with_prefix():
+    f"""long docstring..............................................................."""
+
+
+def multiline_docstring_at_line_limit():
+    """first line-----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""
+
+
+def multiline_docstring_at_line_limit_with_prefix():
+    f"""first line----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""
+
+
+# output
+
+
+def docstring_almost_at_line_limit():
+    """long docstring.................................................................
+    """
+
+
+def docstring_almost_at_line_limit_with_prefix():
+    f"""long docstring................................................................
+    """
+
+
+def mulitline_docstring_almost_at_line_limit():
+    """long docstring.................................................................
+
+    ..................................................................................
+    """
+
+
+def mulitline_docstring_almost_at_line_limit_with_prefix():
+    f"""long docstring................................................................
+
+    ..................................................................................
+    """
+
+
+def docstring_at_line_limit():
+    """long docstring................................................................"""
+
+
+def docstring_at_line_limit_with_prefix():
+    f"""long docstring..............................................................."""
+
+
+def multiline_docstring_at_line_limit():
+    """first line-----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""
+
+
+def multiline_docstring_at_line_limit_with_prefix():
+    f"""first line----------------------------------------------------------------------
+
+    second line----------------------------------------------------------------------"""

tests/test_format.py

@@ -91,6 +91,7 @@
     "one_element_subscript",
     "remove_await_parens",
     "return_annotation_brackets",
+    "docstring_preview",
 ]
 
 SOURCES: List[str] = [