Preserve triple quotes and prefixes for strings

crates/red_knot_python_semantic/src/types/display.rs

@@ -3,7 +3,7 @@
 use std::fmt::{self, Display, Formatter, Write};
 
 use ruff_db::display::FormatterJoinExtension;
-use ruff_python_ast::str::Quote;
+use ruff_python_ast::str::{Quote, TripleQuotes};
 use ruff_python_literal::escape::AsciiEscape;
 
 use crate::types::class_base::ClassBase;
@@ -98,7 +98,7 @@ impl Display for DisplayRepresentation<'_> {
                 let escape =
                     AsciiEscape::with_preferred_quote(bytes.value(self.db).as_ref(), Quote::Double);
 
-                escape.bytes_repr().write(f)
+                escape.bytes_repr(TripleQuotes::No).write(f)
             }
             Type::SliceLiteral(slice) => {
                 f.write_str("slice[")?;

crates/ruff_linter/src/rules/flake8_simplify/rules/split_static_string.rs

@@ -3,8 +3,8 @@ use std::cmp::Ordering;
 use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::{
-    Expr, ExprCall, ExprContext, ExprList, ExprUnaryOp, StringLiteral, StringLiteralFlags,
-    StringLiteralValue, UnaryOp,
+    str::TripleQuotes, Expr, ExprCall, ExprContext, ExprList, ExprUnaryOp, StringLiteral,
+    StringLiteralFlags, StringLiteralValue, UnaryOp,
 };
 use ruff_text_size::{Ranged, TextRange};
 
@@ -123,7 +123,17 @@ fn construct_replacement(elts: &[&str], flags: StringLiteralFlags) -> Expr {
                 Expr::from(StringLiteral {
                     value: Box::from(*elt),
                     range: TextRange::default(),
-                    flags,
+                    // intentionally omit the triple quote flag, if set, to avoid strange
+                    // replacements like
+                    //
+                    // ```python
+                    // """
+                    // itemA
+                    // itemB
+                    // itemC
+                    // """.split() # -> ["""itemA""", """itemB""", """itemC"""]
+                    // ```
+                    flags: flags.with_triple_quotes_set_to(TripleQuotes::No),
                 })
             })
             .collect(),

crates/ruff_python_ast/src/nodes.rs

@@ -14,6 +14,7 @@ use itertools::Itertools;
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::name::Name;
+use crate::str::TripleQuotes;
 use crate::{
     int,
     str::Quote,
@@ -981,25 +982,24 @@ pub trait StringFlags: Copy {
     /// Does the string use single or double quotes in its opener and closer?
     fn quote_style(self) -> Quote;
 
-    /// Is the string triple-quoted, i.e.,
-    /// does it begin and end with three consecutive quote characters?
-    fn is_triple_quoted(self) -> bool;
+    fn triple_quotes(self) -> TripleQuotes;
 
     fn prefix(self) -> AnyStringPrefix;
 
+    /// Is the string triple-quoted, i.e.,
+    /// does it begin and end with three consecutive quote characters?
+    fn is_triple_quoted(self) -> bool {
+        self.triple_quotes().is_yes()
+    }
+
     /// A `str` representation of the quotes used to start and close.
     /// This does not include any prefixes the string has in its opener.
     fn quote_str(self) -> &'static str {
-        if self.is_triple_quoted() {
-            match self.quote_style() {
-                Quote::Single => "'''",
-                Quote::Double => r#"""""#,
-            }
-        } else {
-            match self.quote_style() {
-                Quote::Single => "'",
-                Quote::Double => "\"",
-            }
+        match (self.triple_quotes(), self.quote_style()) {
+            (TripleQuotes::Yes, Quote::Single) => "'''",
+            (TripleQuotes::Yes, Quote::Double) => r#"""""#,
+            (TripleQuotes::No, Quote::Single) => "'",
+            (TripleQuotes::No, Quote::Double) => "\"",
         }
     }
 
@@ -1029,9 +1029,18 @@ pub trait StringFlags: Copy {
     }
 
     fn format_string_contents(self, contents: &str) -> String {
-        let prefix = self.prefix();
+        let buf_size = self.opener_len().to_usize() + contents.len() + self.closer_len().to_usize();
+        let mut buffer = String::with_capacity(buf_size);
+        self.write_string_contents(&mut buffer, contents);
+        buffer
+    }
+
+    fn write_string_contents(self, buffer: &mut String, contents: &str) {
         let quote_str = self.quote_str();
-        format!("{prefix}{quote_str}{contents}{quote_str}")
+        buffer.push_str(self.prefix().as_str());
+        buffer.push_str(quote_str);
+        buffer.push_str(contents);
+        buffer.push_str(quote_str);
     }
 }
 
@@ -1078,7 +1087,7 @@ pub struct FStringFlags(FStringFlagsInner);
 impl FStringFlags {
     /// Construct a new [`FStringFlags`] with **no flags set**.
     ///
-    /// See [`FStringFlags::with_quote_style`], [`FStringFlags::with_triple_quotes`], and
+    /// See [`FStringFlags::with_quote_style`], [`FStringFlags::with_triple_quotes_set_to`], and
     /// [`FStringFlags::with_prefix`] for ways of setting the quote style (single or double),
     /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
     ///
@@ -1097,8 +1106,9 @@ impl FStringFlags {
     }
 
     #[must_use]
-    pub fn with_triple_quotes(mut self) -> Self {
-        self.0 |= FStringFlagsInner::TRIPLE_QUOTED;
+    pub fn with_triple_quotes_set_to(mut self, triple_quotes: TripleQuotes) -> Self {
+        self.0
+            .set(FStringFlagsInner::TRIPLE_QUOTED, triple_quotes.is_yes());
         self
     }
 
@@ -1132,8 +1142,12 @@ impl StringFlags for FStringFlags {
     /// Return `true` if the f-string is triple-quoted, i.e.,
     /// it begins and ends with three consecutive quote characters.
     /// For example: `f"""{bar}"""`
-    fn is_triple_quoted(self) -> bool {
-        self.0.contains(FStringFlagsInner::TRIPLE_QUOTED)
+    fn triple_quotes(self) -> TripleQuotes {
+        if self.0.contains(FStringFlagsInner::TRIPLE_QUOTED) {
+            TripleQuotes::Yes
+        } else {
+            TripleQuotes::No
+        }
     }
 
     /// Return the quoting style (single or double quotes)
@@ -1458,9 +1472,10 @@ pub struct StringLiteralFlags(StringLiteralFlagsInner);
 impl StringLiteralFlags {
     /// Construct a new [`StringLiteralFlags`] with **no flags set**.
     ///
-    /// See [`StringLiteralFlags::with_quote_style`], [`StringLiteralFlags::with_triple_quotes`],
-    /// and [`StringLiteralFlags::with_prefix`] for ways of setting the quote style (single or
-    /// double), enabling triple quotes, and adding prefixes (such as `r` or `u`), respectively.
+    /// See [`StringLiteralFlags::with_quote_style`],
+    /// [`StringLiteralFlags::with_triple_quotes_set_to`], and [`StringLiteralFlags::with_prefix`]
+    /// for ways of setting the quote style (single or double), enabling triple quotes, and adding
+    /// prefixes (such as `r` or `u`), respectively.
     ///
     /// See the documentation for [`StringLiteralFlags`] for additional caveats on this constructor,
     /// and situations in which alternative ways to construct this struct should be used, especially
@@ -1477,8 +1492,11 @@ impl StringLiteralFlags {
     }
 
     #[must_use]
-    pub fn with_triple_quotes(mut self) -> Self {
-        self.0 |= StringLiteralFlagsInner::TRIPLE_QUOTED;
+    pub fn with_triple_quotes_set_to(mut self, triple_quotes: TripleQuotes) -> Self {
+        self.0.set(
+            StringLiteralFlagsInner::TRIPLE_QUOTED,
+            triple_quotes.is_yes(),
+        );
         self
     }
 
@@ -1550,8 +1568,12 @@ impl StringFlags for StringLiteralFlags {
     /// Return `true` if the string is triple-quoted, i.e.,
     /// it begins and ends with three consecutive quote characters.
     /// For example: `"""bar"""`
-    fn is_triple_quoted(self) -> bool {
-        self.0.contains(StringLiteralFlagsInner::TRIPLE_QUOTED)
+    fn triple_quotes(self) -> TripleQuotes {
+        if self.0.contains(StringLiteralFlagsInner::TRIPLE_QUOTED) {
+            TripleQuotes::Yes
+        } else {
+            TripleQuotes::No
+        }
     }
 
     fn prefix(self) -> AnyStringPrefix {
@@ -1847,9 +1869,10 @@ pub struct BytesLiteralFlags(BytesLiteralFlagsInner);
 impl BytesLiteralFlags {
     /// Construct a new [`BytesLiteralFlags`] with **no flags set**.
     ///
-    /// See [`BytesLiteralFlags::with_quote_style`], [`BytesLiteralFlags::with_triple_quotes`], and
-    /// [`BytesLiteralFlags::with_prefix`] for ways of setting the quote style (single or double),
-    /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
+    /// See [`BytesLiteralFlags::with_quote_style`],
+    /// [`BytesLiteralFlags::with_triple_quotes_set_to`], and [`BytesLiteralFlags::with_prefix`] for
+    /// ways of setting the quote style (single or double), enabling triple quotes, and adding
+    /// prefixes (such as `r`), respectively.
     ///
     /// See the documentation for [`BytesLiteralFlags`] for additional caveats on this constructor,
     /// and situations in which alternative ways to construct this struct should be used, especially
@@ -1866,8 +1889,11 @@ impl BytesLiteralFlags {
     }
 
     #[must_use]
-    pub fn with_triple_quotes(mut self) -> Self {
-        self.0 |= BytesLiteralFlagsInner::TRIPLE_QUOTED;
+    pub fn with_triple_quotes_set_to(mut self, triple_quotes: TripleQuotes) -> Self {
+        self.0.set(
+            BytesLiteralFlagsInner::TRIPLE_QUOTED,
+            triple_quotes.is_yes(),
+        );
         self
     }
 
@@ -1910,8 +1936,12 @@ impl StringFlags for BytesLiteralFlags {
     /// Return `true` if the bytestring is triple-quoted, i.e.,
     /// it begins and ends with three consecutive quote characters.
     /// For example: `b"""{bar}"""`
-    fn is_triple_quoted(self) -> bool {
-        self.0.contains(BytesLiteralFlagsInner::TRIPLE_QUOTED)
+    fn triple_quotes(self) -> TripleQuotes {
+        if self.0.contains(BytesLiteralFlagsInner::TRIPLE_QUOTED) {
+            TripleQuotes::Yes
+        } else {
+            TripleQuotes::No
+        }
     }
 
     /// Return the quoting style (single or double quotes)
@@ -2035,7 +2065,7 @@ bitflags! {
     }
 }
 
-#[derive(Default, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct AnyStringFlags(AnyStringFlagsInner);
 
 impl AnyStringFlags {
@@ -2073,13 +2103,11 @@ impl AnyStringFlags {
         self
     }
 
-    pub fn new(prefix: AnyStringPrefix, quotes: Quote, triple_quoted: bool) -> Self {
-        let new = Self::default().with_prefix(prefix).with_quote_style(quotes);
-        if triple_quoted {
-            new.with_triple_quotes()
-        } else {
-            new
-        }
+    pub fn new(prefix: AnyStringPrefix, quotes: Quote, triple_quotes: TripleQuotes) -> Self {
+        Self(AnyStringFlagsInner::empty())
+            .with_prefix(prefix)
+            .with_quote_style(quotes)
+            .with_triple_quotes_set_to(triple_quotes)
     }
 
     /// Does the string have a `u` or `U` prefix?
@@ -2114,8 +2142,9 @@ impl AnyStringFlags {
     }
 
     #[must_use]
-    pub fn with_triple_quotes(mut self) -> Self {
-        self.0 |= AnyStringFlagsInner::TRIPLE_QUOTED;
+    pub fn with_triple_quotes_set_to(mut self, triple_quotes: TripleQuotes) -> Self {
+        self.0
+            .set(AnyStringFlagsInner::TRIPLE_QUOTED, triple_quotes.is_yes());
         self
     }
 }
@@ -2130,10 +2159,12 @@ impl StringFlags for AnyStringFlags {
         }
     }
 
-    /// Is the string triple-quoted, i.e.,
-    /// does it begin and end with three consecutive quote characters?
-    fn is_triple_quoted(self) -> bool {
-        self.0.contains(AnyStringFlagsInner::TRIPLE_QUOTED)
+    fn triple_quotes(self) -> TripleQuotes {
+        if self.0.contains(AnyStringFlagsInner::TRIPLE_QUOTED) {
+            TripleQuotes::Yes
+        } else {
+            TripleQuotes::No
+        }
     }
 
     fn prefix(self) -> AnyStringPrefix {
@@ -2193,14 +2224,10 @@ impl From<AnyStringFlags> for StringLiteralFlags {
                 value.prefix()
             )
         };
-        let new = StringLiteralFlags::empty()
+        StringLiteralFlags::empty()
             .with_quote_style(value.quote_style())
-            .with_prefix(prefix);
-        if value.is_triple_quoted() {
-            new.with_triple_quotes()
-        } else {
-            new
-        }
+            .with_prefix(prefix)
+            .with_triple_quotes_set_to(value.triple_quotes())
     }
 }
 
@@ -2209,7 +2236,7 @@ impl From<StringLiteralFlags> for AnyStringFlags {
         Self::new(
             AnyStringPrefix::Regular(value.prefix()),
             value.quote_style(),
-            value.is_triple_quoted(),
+            value.triple_quotes(),
         )
     }
 }
@@ -2222,14 +2249,10 @@ impl From<AnyStringFlags> for BytesLiteralFlags {
                 value.prefix()
             )
         };
-        let new = BytesLiteralFlags::empty()
+        BytesLiteralFlags::empty()
             .with_quote_style(value.quote_style())
-            .with_prefix(bytestring_prefix);
-        if value.is_triple_quoted() {
-            new.with_triple_quotes()
-        } else {
-            new
-        }
+            .with_prefix(bytestring_prefix)
+            .with_triple_quotes_set_to(value.triple_quotes())
     }
 }
 
@@ -2238,7 +2261,7 @@ impl From<BytesLiteralFlags> for AnyStringFlags {
         Self::new(
             AnyStringPrefix::Bytes(value.prefix()),
             value.quote_style(),
-            value.is_triple_quoted(),
+            value.triple_quotes(),
         )
     }
 }
@@ -2251,14 +2274,10 @@ impl From<AnyStringFlags> for FStringFlags {
                 value.prefix()
             )
         };
-        let new = FStringFlags::empty()
+        FStringFlags::empty()
             .with_quote_style(value.quote_style())
-            .with_prefix(fstring_prefix);
-        if value.is_triple_quoted() {
-            new.with_triple_quotes()
-        } else {
-            new
-        }
+            .with_prefix(fstring_prefix)
+            .with_triple_quotes_set_to(value.triple_quotes())
     }
 }
 
@@ -2267,7 +2286,7 @@ impl From<FStringFlags> for AnyStringFlags {
         Self::new(
             AnyStringPrefix::Format(value.prefix()),
             value.quote_style(),
-            value.is_triple_quoted(),
+            value.triple_quotes(),
         )
     }
 }