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

idorrington92 · felix-hilden · web-flow · commit 20d8ccb54253 · 2022-05-07T21:34:28.000-07:00
Fixes psf#1632 Co-authored-by: Felix Hildén <felix.hilden@gmail.com>
diff --git a/CHANGES.md b/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)
diff --git a/src/black/linegen.py b/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,7 +329,29 @@ def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
 
             # We could enforce triple quotes at this point.
             quote = quote_char * quote_len
-            leaf.value = prefix + quote + docstring + quote
+
+            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)
 
diff --git a/src/black/mode.py b/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):
diff --git a/tests/data/docstring.py b/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----------------------------------------------------------------------"""
diff --git a/tests/data/docstring_preview.py b/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----------------------------------------------------------------------"""
diff --git a/tests/test_format.py b/tests/test_format.py
@@ -91,6 +91,7 @@
     "one_element_subscript",
     "remove_await_parens",
     "return_annotation_brackets",
+    "docstring_preview",
 ]
 
 SOURCES: List[str] = [

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,7 @@`
`91`	`91`	`"one_element_subscript",`
`92`	`92`	`"remove_await_parens",`
`93`	`93`	`"return_annotation_brackets",`
	`94`	`+ "docstring_preview",`
`94`	`95`	`]`
`95`	`96`
`96`	`97`	`SOURCES: List[str] = [`