From 9270a10f6f59f069eb14ffba0c75f58e5895b27c Mon Sep 17 00:00:00 2001
From: =?utf8?q?=C5=81ukasz=20Langa?= <lukasz@langa.pl>
Date: Tue, 25 Aug 2020 22:26:13 +0200
Subject: [PATCH] Improve docstring re-indentation handling

This addresses a few crashers, namely:

* producing non-equivalent code due to mangling escaped newlines,

* invalid hugging quote characters in the docstring body to the docstring outer
  triple quotes (causing a quadruple quote which is a syntax error),

* lack of handling for docstrings that start on the same line as the `def`, and

* invalid stripping of outer triple quotes when the docstring contained
  a string prefix.

As a bonus, tests now also run when string normalization is disabled.
---
 src/black/__init__.py   | 42 +++++++++++++++++++++++------
 tests/data/docstring.py | 59 +++++++++++++++++++++++++++++++++++++++++
 tests/test_black.py     |  5 ++++
 3 files changed, 98 insertions(+), 8 deletions(-)

diff --git a/src/black/__init__.py b/src/black/__init__.py
index e37caa9..c3c8c20 100644
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -2037,13 +2037,20 @@ class LineGenerator(Visitor[Line]):
         yield from self.visit_default(node)
 
     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
-        # Check if it's a docstring
-        if prev_siblings_are(
-            leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
-        ) and is_multiline_string(leaf):
-            prefix = "    " * self.current_line.depth
-            docstring = fix_docstring(leaf.value[3:-3], prefix)
-            leaf.value = leaf.value[0:3] + docstring + leaf.value[-3:]
+        if is_docstring(leaf) and "\\\n" not in leaf.value:
+            # We're ignoring docstrings with backslash newline escapes because changing
+            # indentation of those changes the AST representation of the code.
+            prefix = get_string_prefix(leaf.value)
+            lead_len = len(prefix) + 3
+            tail_len = -3
+            indent = " " * 4 * self.current_line.depth
+            docstring = fix_docstring(leaf.value[lead_len:tail_len], indent)
+            if docstring:
+                if leaf.value[lead_len - 1] == docstring[0]:
+                    docstring = " " + docstring
+                if leaf.value[tail_len + 1] == docstring[-1]:
+                    docstring = docstring + " "
+            leaf.value = leaf.value[0:lead_len] + docstring + leaf.value[tail_len:]
             normalize_string_quotes(leaf)
 
         yield from self.visit_default(leaf)
@@ -6608,6 +6615,26 @@ def patched_main() -> None:
     main()
 
 
+def is_docstring(leaf: Leaf) -> bool:
+    if not is_multiline_string(leaf):
+        # For the purposes of docstring re-indentation, we don't need to do anything
+        # with single-line docstrings.
+        return False
+
+    if prev_siblings_are(
+        leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
+    ):
+        return True
+
+    # Multiline docstring on the same line as the `def`.
+    if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
+        # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
+        # grammar. We're safe to return True without further checks.
+        return True
+
+    return False
+
+
 def fix_docstring(docstring: str, prefix: str) -> str:
     # https://www.python.org/dev/peps/pep-0257/#handling-docstring-indentation
     if not docstring:
@@ -6631,7 +6658,6 @@ def fix_docstring(docstring: str, prefix: str) -> str:
                 trimmed.append(prefix + stripped_line)
             else:
                 trimmed.append("")
-    # Return a single string:
     return "\n".join(trimmed)
 
 
diff --git a/tests/data/docstring.py b/tests/data/docstring.py
index fcb8eb1..2d3d73a 100644
--- a/tests/data/docstring.py
+++ b/tests/data/docstring.py
@@ -81,6 +81,35 @@ def single_line():
     """
     pass
 
+
+def this():
+    r"""
+    'hey ho'
+    """
+
+
+def that():
+  """ "hey yah" """
+
+
+def and_that():
+  """
+  "hey yah" """
+
+
+def and_this():
+  ''' 
+  "hey yah"'''
+
+
+def believe_it_or_not_this_is_in_the_py_stdlib(): ''' 
+"hey yah"'''
+
+
+def ignored_docstring():
+    """a => \
+b"""  
+
 # output
 
 class MyClass:
@@ -164,3 +193,33 @@ def over_indent():
 def single_line():
     """But with a newline after it!"""
     pass
+
+
+def this():
+    r"""
+    'hey ho'
+    """
+
+
+def that():
+    """ "hey yah" """
+
+
+def and_that():
+    """
+    "hey yah" """
+
+
+def and_this():
+    '''
+    "hey yah"'''
+
+
+def believe_it_or_not_this_is_in_the_py_stdlib():
+    '''
+    "hey yah"'''
+
+
+def ignored_docstring():
+    """a => \
+b"""
\ No newline at end of file
diff --git a/tests/test_black.py b/tests/test_black.py
index 6705490..cf311f5 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -496,6 +496,11 @@ class BlackTestCase(unittest.TestCase):
         self.assertFormatEqual(expected, actual)
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, DEFAULT_MODE)
+        mode = replace(DEFAULT_MODE, string_normalization=False)
+        not_normalized = fs(source, mode=mode)
+        self.assertFormatEqual(expected, not_normalized)
+        black.assert_equivalent(source, not_normalized)
+        black.assert_stable(source, not_normalized, mode=mode)
 
     def test_long_strings(self) -> None:
         """Tests for splitting long strings."""
-- 
2.39.2