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>

diff --git a/CHANGES.md b/CHANGES.md
index 4cea7fceaad783830743a4e96cd22f419933e092..8f43431c8421c23d9986697a1f1295fd848f06d4 100644 (file)
--- 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
index 91fdeef8f2f94a3e1b7b8cb244b19723beac9b85..ff54e05c4e6299d1bd2a456473a169839a03b43a 100644 (file)
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -305,9 +305,9 @@ class LineGenerator(Visitor[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 @@ class LineGenerator(Visitor[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
index 6bd4ce1442118d7f10e661bae25db9bdcfb7a41b..a418e0eb66510ef6ea451cbc51f8b4be035816d4 100644 (file)
--- 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
index 96bcf525b16fe56538739f4c895e021a4766dadd..7153be468c112963806ea2a7623717b1c3a88b7b 100644 (file)
--- 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
new file mode 100644 (file)
index 0000000..2da4cd1
--- /dev/null
+++ 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
index 1916146e84dc570a36a38a0e940e0c207c5b6ce2..2f08d1f273d4040c4f62b09cba3bb0fea0b80a71 100644 (file)
--- a/tests/test_format.py
+++ b/tests/test_format.py
@@ -91,6 +91,7 @@ PREVIEW_CASES: List[str] = [
     "one_element_subscript",
     "remove_await_parens",
     "return_annotation_brackets",
+    "docstring_preview",
 ]
 
 SOURCES: List[str] = [