Improve docstring re-indentation handling

src/black/__init__.py

@@ -2037,13 +2037,20 @@ def visit_factor(self, node: Node) -> Iterator[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)
 
 

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"""

tests/test_black.py

@@ -496,6 +496,11 @@ def test_docstring(self) -> None:
         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."""