Don't add chaperone space after escaped quote in triple quote (#17216)

Co-authored-by: Micha Reiser <micha@reiser.io>

Author: maxmynter

crates/ruff_python_formatter/resources/test/fixtures/ruff/docstring_chaperones.py

@@ -0,0 +1,66 @@
+def a1():
+    """Needs chaperone\\" """
+
+def a2():
+    """Needs chaperone\\\ """
+
+def a3():
+    """Needs chaperone" """
+
+def a4():
+    """Needs chaperone\ """
+
+def a5():
+    """Needs chaperone\\\\\ """
+
+def a6():
+    """Needs chaperone\"" """
+
+def a7():
+    """Doesn't need chaperone\""""
+
+def a8():
+    """Doesn't need chaperone\'"""
+
+def a9():
+    """Doesn't need chaperone\\\""""
+
+def a10():
+    """Doesn't need chaperone\\\'"""
+
+def a11():
+    """Doesn't need chaperone; all slashes escaped\\\\"""
+
+def a12():
+    """Doesn't need chaperone\\"""
+
+def a12():
+    """Doesn't need "chaperone" with contained quotes"""
+
+def a13():
+    """Doesn't need chaperone\\\"\"\""""
+
+def a14():
+    """Multiline docstring
+    doesn't need chaperone
+    """
+
+def a15():
+    """Multiline docstring
+    doesn't need chaperone\
+    """
+
+def a16():
+    """Multiline docstring with
+    odd number of backslashes don't need chaperone\\\
+    """
+
+def a17():
+    """Multiline docstring with
+    even number of backslashes don't need chaperone\\\\
+    """
+
+def a18():
+    r"""Raw multiline docstring
+    doesn't need chaperone\
+    """

crates/ruff_python_formatter/src/preview.rs

@@ -13,3 +13,10 @@ pub(crate) const fn is_hug_parens_with_braces_and_square_brackets_enabled(
 ) -> bool {
     context.is_preview()
 }
+
+/// Returns `true` if the [`no_chaperone_for_escaped_quote_in_triple_quoted_docstring`](https://github.com/astral-sh/ruff/pull/17216) preview style is enabled.
+pub(crate) const fn is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled(
+    context: &PyFormatContext,
+) -> bool {
+    context.is_preview()
+}

crates/ruff_python_formatter/src/string/docstring.rs

@@ -19,11 +19,11 @@ use {
     ruff_text_size::{Ranged, TextLen, TextRange, TextSize},
 };
 
+use super::NormalizedString;
+use crate::preview::is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled;
 use crate::string::StringQuotes;
 use crate::{prelude::*, DocstringCodeLineWidth, FormatModuleError};
 
-use super::NormalizedString;
-
 /// Format a docstring by trimming whitespace and adjusting the indentation.
 ///
 /// Summary of changes we make:
@@ -168,7 +168,7 @@ pub(crate) fn format(normalized: &NormalizedString, f: &mut PyFormatter) -> Form
     if docstring[first.len()..].trim().is_empty() {
         // For `"""\n"""` or other whitespace between the quotes, black keeps a single whitespace,
         // but `""""""` doesn't get one inserted.
-        if needs_chaperone_space(normalized.flags(), trim_end)
+        if needs_chaperone_space(normalized.flags(), trim_end, f.context())
             || (trim_end.is_empty() && !docstring.is_empty())
         {
             space().fmt(f)?;
@@ -208,7 +208,7 @@ pub(crate) fn format(normalized: &NormalizedString, f: &mut PyFormatter) -> Form
     let trim_end = docstring
         .as_ref()
         .trim_end_matches(|c: char| c.is_whitespace() && c != '\n');
-    if needs_chaperone_space(normalized.flags(), trim_end) {
+    if needs_chaperone_space(normalized.flags(), trim_end, f.context()) {
         space().fmt(f)?;
     }
 
@@ -1596,17 +1596,45 @@ fn docstring_format_source(
     Ok(formatted.print()?)
 }
 
-/// If the last line of the docstring is `content" """` or `content\ """`, we need a chaperone space
-/// that avoids `content""""` and `content\"""`. This does only applies to un-escaped backslashes,
-/// so `content\\ """` doesn't need a space while `content\\\ """` does.
-pub(super) fn needs_chaperone_space(flags: AnyStringFlags, trim_end: &str) -> bool {
-    if trim_end.chars().rev().take_while(|c| *c == '\\').count() % 2 == 1 {
-        true
+/// If the last line of the docstring is `content""""` or `content\"""`, we need a chaperone space
+/// that avoids `content""""` and `content\"""`. This only applies to un-escaped backslashes,
+/// so `content\\"""` doesn't need a space while `content\\\"""` does.
+pub(super) fn needs_chaperone_space(
+    flags: AnyStringFlags,
+    trim_end: &str,
+    context: &PyFormatContext,
+) -> bool {
+    if count_consecutive_chars_from_end(trim_end, '\\') % 2 == 1 {
+        // Odd backslash count; chaperone avoids escaping closing quotes
+        // `"\ "` -> prevent that this becomes `"\"` which escapes the closing quote.
+        return true;
+    }
+
+    if is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled(context) {
+        if flags.is_triple_quoted() {
+            if let Some(before_quote) = trim_end.strip_suffix(flags.quote_style().as_char()) {
+                if count_consecutive_chars_from_end(before_quote, '\\') % 2 == 0 {
+                    // Even backslash count preceding quote;
+                    // ```py
+                    // """a "  """
+                    // """a \\"  """
+                    // ```
+                    // The chaperon is needed or the triple quoted string "ends" with 4 instead of 3 quotes.
+                    return true;
+                }
+            }
+        }
+
+        false
     } else {
         flags.is_triple_quoted() && trim_end.ends_with(flags.quote_style().as_char())
     }
 }
 
+fn count_consecutive_chars_from_end(s: &str, target: char) -> usize {
+    s.chars().rev().take_while(|c| *c == target).count()
+}
+
 #[derive(Copy, Clone, Debug)]
 enum Indentation {
     /// Space only indentation or an empty indentation.

crates/ruff_python_formatter/src/string/implicit.rs

@@ -372,7 +372,7 @@ impl Format<PyFormatContext<'_>> for FormatLiteralContent {
                 Cow::Owned(normalized) => text(normalized).fmt(f)?,
             }
 
-            if self.trim_end && needs_chaperone_space(self.flags, &normalized) {
+            if self.trim_end && needs_chaperone_space(self.flags, &normalized, f.context()) {
                 space().fmt(f)?;
             }
         }

crates/ruff_python_formatter/tests/snapshots/format@docstring_chaperones.py.snap

@@ -0,0 +1,195 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/docstring_chaperones.py
+---
+## Input
+```python
+def a1():
+    """Needs chaperone\\" """
+
+def a2():
+    """Needs chaperone\\\ """
+
+def a3():
+    """Needs chaperone" """
+
+def a4():
+    """Needs chaperone\ """
+
+def a5():
+    """Needs chaperone\\\\\ """
+
+def a6():
+    """Needs chaperone\"" """
+
+def a7():
+    """Doesn't need chaperone\""""
+
+def a8():
+    """Doesn't need chaperone\'"""
+
+def a9():
+    """Doesn't need chaperone\\\""""
+
+def a10():
+    """Doesn't need chaperone\\\'"""
+
+def a11():
+    """Doesn't need chaperone; all slashes escaped\\\\"""
+
+def a12():
+    """Doesn't need chaperone\\"""
+
+def a12():
+    """Doesn't need "chaperone" with contained quotes"""
+
+def a13():
+    """Doesn't need chaperone\\\"\"\""""
+
+def a14():
+    """Multiline docstring
+    doesn't need chaperone
+    """
+
+def a15():
+    """Multiline docstring
+    doesn't need chaperone\
+    """
+
+def a16():
+    """Multiline docstring with
+    odd number of backslashes don't need chaperone\\\
+    """
+
+def a17():
+    """Multiline docstring with
+    even number of backslashes don't need chaperone\\\\
+    """
+
+def a18():
+    r"""Raw multiline docstring
+    doesn't need chaperone\
+    """
+```
+
+## Output
+```python
+def a1():
+    """Needs chaperone\\" """
+
+
+def a2():
+    """Needs chaperone\\\ """
+
+
+def a3():
+    """Needs chaperone" """
+
+
+def a4():
+    """Needs chaperone\ """
+
+
+def a5():
+    """Needs chaperone\\\\\ """
+
+
+def a6():
+    """Needs chaperone\"" """
+
+
+def a7():
+    """Doesn't need chaperone\" """
+
+
+def a8():
+    """Doesn't need chaperone\'"""
+
+
+def a9():
+    """Doesn't need chaperone\\\" """
+
+
+def a10():
+    """Doesn't need chaperone\\\'"""
+
+
+def a11():
+    """Doesn't need chaperone; all slashes escaped\\\\"""
+
+
+def a12():
+    """Doesn't need chaperone\\"""
+
+
+def a12():
+    """Doesn't need "chaperone" with contained quotes"""
+
+
+def a13():
+    """Doesn't need chaperone\\\"\"\" """
+
+
+def a14():
+    """Multiline docstring
+    doesn't need chaperone
+    """
+
+
+def a15():
+    """Multiline docstring
+    doesn't need chaperone\
+    """
+
+
+def a16():
+    """Multiline docstring with
+    odd number of backslashes don't need chaperone\\\
+    """
+
+
+def a17():
+    """Multiline docstring with
+    even number of backslashes don't need chaperone\\\\
+    """
+
+
+def a18():
+    r"""Raw multiline docstring
+    doesn't need chaperone\
+    """
+```
+
+
+## Preview changes
+```diff
+--- Stable
++++ Preview
+@@ -23,7 +23,7 @@
+ 
+ 
+ def a7():
+-    """Doesn't need chaperone\" """
++    """Doesn't need chaperone\""""
+ 
+ 
+ def a8():
+@@ -31,7 +31,7 @@
+ 
+ 
+ def a9():
+-    """Doesn't need chaperone\\\" """
++    """Doesn't need chaperone\\\""""
+ 
+ 
+ def a10():
+@@ -51,7 +51,7 @@
+ 
+ 
+ def a13():
+-    """Doesn't need chaperone\\\"\"\" """
++    """Doesn't need chaperone\\\"\"\""""
+ 
+ 
+ def a14():
+```