Unify StringPart and AnyStringPart

MichaReiser · MichaReiser · commit 22cb7d4e7f90 · 2024-10-15T11:30:26.000+02:00
diff --git a/crates/ruff_python_formatter/src/expression/expr_f_string.rs b/crates/ruff_python_formatter/src/expression/expr_f_string.rs
@@ -17,11 +17,10 @@ impl FormatNodeRule<ExprFString> for FormatExprFString {
         let ExprFString { value, .. } = item;
 
         match value.as_slice() {
-            [f_string_part] => FormatFStringPart::new(
-                f_string_part,
-                f_string_quoting(item, f.context().locator()),
-            )
-            .fmt(f),
+            [f_string_part] => {
+                FormatFStringPart::new(f_string_part, f_string_quoting(item, f.context().locator()))
+                    .fmt(f)
+            }
             _ => in_parentheses_only_group(&FormatImplicitConcatenatedString::new(item)).fmt(f),
         }
     }
diff --git a/crates/ruff_python_formatter/src/string/any.rs b/crates/ruff_python_formatter/src/string/any.rs
@@ -165,8 +165,8 @@ impl FusedIterator for AnyStringPartsIter<'_> {}
 /// string. This could be either a string, bytes or f-string.
 ///
 /// This is constructed from the [`AnyString::parts`] method on [`AnyString`].
-#[derive(Clone, Debug)]
-pub(super) enum AnyStringPart<'a> {
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum AnyStringPart<'a> {
     String(&'a ast::StringLiteral),
     Bytes(&'a ast::BytesLiteral),
     FString(&'a ast::FString),
@@ -180,6 +180,15 @@ impl AnyStringPart<'_> {
             Self::FString(part) => part.flags.into(),
         }
     }
+
+    /// Returns the range of the string's content in the source (minus prefix and quotes).
+    pub(super) fn content_range(self) -> TextRange {
+        let kind = self.flags();
+        TextRange::new(
+            self.start() + kind.opener_len(),
+            self.end() - kind.closer_len(),
+        )
+    }
 }
 
 impl<'a> From<&AnyStringPart<'a>> for AnyNodeRef<'a> {
@@ -201,3 +210,21 @@ impl Ranged for AnyStringPart<'_> {
         }
     }
 }
+
+impl<'a> From<&'a ast::StringLiteral> for AnyStringPart<'a> {
+    fn from(value: &'a ast::StringLiteral) -> Self {
+        Self::String(value)
+    }
+}
+
+impl<'a> From<&'a ast::BytesLiteral> for AnyStringPart<'a> {
+    fn from(value: &'a ast::BytesLiteral) -> Self {
+        Self::Bytes(value)
+    }
+}
+
+impl<'a> From<&'a ast::FString> for AnyStringPart<'a> {
+    fn from(value: &'a ast::FString) -> Self {
+        Self::FString(value)
+    }
+}
diff --git a/crates/ruff_python_formatter/src/string/mod.rs b/crates/ruff_python_formatter/src/string/mod.rs
@@ -3,11 +3,10 @@ pub(crate) use normalize::{normalize_string, NormalizedString, StringNormalizer}
 use ruff_formatter::format_args;
 use ruff_python_ast::str::Quote;
 use ruff_python_ast::{
-    self as ast,
     str_prefix::{AnyStringPrefix, StringLiteralPrefix},
     AnyStringFlags, StringFlags,
 };
-use ruff_text_size::{Ranged, TextRange};
+use ruff_text_size::Ranged;
 
 use crate::comments::{leading_comments, trailing_comments};
 use crate::expression::parentheses::in_parentheses_only_soft_line_break_or_space;
@@ -42,7 +41,7 @@ impl<'a> FormatImplicitConcatenatedString<'a> {
         }
     }
 
-    fn merged_prefix(&self, context: &PyFormatContext) -> Option<AnyStringPrefix> {
+    fn merged_prefix(&self, context: &PyFormatContext) -> Option<Quote> {
         // Early exit if it's known that this string can't be joined to avoid
         // unnecessary computation of the quote metadata.
         if self.string.parts().any(|part| {
@@ -75,6 +74,7 @@ impl<'a> FormatImplicitConcatenatedString<'a> {
         }
 
         // The string is either a regular string, f-string, or bytes string.
+        let normalizer = StringNormalizer::from_context(context);
 
         // TODO: Do we need to respect the quoting?
         let mut merged_quotes: Option<QuoteMetadata> = None;
@@ -83,28 +83,24 @@ impl<'a> FormatImplicitConcatenatedString<'a> {
         // Possibly run directly on entire string?
 
         for part in self.string.parts() {
-            let flags = part.flags();
-
-            let comments = context.comments().leading_dangling_trailing(&part);
-
-            // Can't merge parts with comments.
-            if comments.has_leading() || comments.has_trailing() {
+            let Ok(preferred_quote) = Quote::try_from(normalizer.preferred_quote_style(part))
+            else {
+                // TOOD: Handle preserve in some way or another
                 return None;
-            }
+            };
 
-            // let part_quote_metadata = QuoteMetadata::from_part(part, context);
+            // Again, this takes a StringPart and not a `AnyStringPart`.
+            let part_quote_metadata = QuoteMetadata::from_part(part, preferred_quote, context);
 
-            // if let Some(merged) = merged_quotes.as_mut() {
-            //     // FIXME: this is not correct.
-            //     *merged = merged.with_prefix(flags.prefix());
-            // } else {
-            //     merged_flags = Some(part.flags());
-            // }
+            if let Some(merged) = merged_quotes.as_mut() {
+                // FIXME: this is not correct.
+                *merged = part_quote_metadata.merge(merged)?;
+            } else {
+                merged_quotes = Some(part_quote_metadata);
+            }
         }
 
-        // merged_flags.map(StringFlags::prefix)
-
-        todo!()
+        Some(merged_quotes?.choose(Quote::Double))
     }
 }
 
@@ -214,58 +210,3 @@ impl From<Quote> for QuoteStyle {
         }
     }
 }
-
-#[derive(Debug, Clone, Copy)]
-pub(crate) struct StringPart {
-    flags: AnyStringFlags,
-    range: TextRange,
-}
-
-impl Ranged for StringPart {
-    fn range(&self) -> TextRange {
-        self.range
-    }
-}
-
-impl StringPart {
-    /// Use the `kind()` method to retrieve information about the
-    fn flags(self) -> AnyStringFlags {
-        self.flags
-    }
-
-    /// Returns the range of the string's content in the source (minus prefix and quotes).
-    fn content_range(self) -> TextRange {
-        let kind = self.flags();
-        TextRange::new(
-            self.start() + kind.opener_len(),
-            self.end() - kind.closer_len(),
-        )
-    }
-}
-
-impl From<&ast::StringLiteral> for StringPart {
-    fn from(value: &ast::StringLiteral) -> Self {
-        Self {
-            range: value.range,
-            flags: value.flags.into(),
-        }
-    }
-}
-
-impl From<&ast::BytesLiteral> for StringPart {
-    fn from(value: &ast::BytesLiteral) -> Self {
-        Self {
-            range: value.range,
-            flags: value.flags.into(),
-        }
-    }
-}
-
-impl From<&ast::FString> for StringPart {
-    fn from(value: &ast::FString) -> Self {
-        Self {
-            range: value.range,
-            flags: value.flags.into(),
-        }
-    }
-}
diff --git a/crates/ruff_python_formatter/src/string/normalize.rs b/crates/ruff_python_formatter/src/string/normalize.rs
@@ -9,7 +9,8 @@ use ruff_text_size::{Ranged, TextRange};
 use crate::context::FStringState;
 use crate::prelude::*;
 use crate::preview::is_f_string_formatting_enabled;
-use crate::string::{Quoting, StringPart, StringQuotes};
+use crate::string::any::AnyStringPart;
+use crate::string::{Quoting, StringQuotes};
 use crate::QuoteStyle;
 
 pub(crate) struct StringNormalizer<'a, 'src> {
@@ -37,7 +38,7 @@ impl<'a, 'src> StringNormalizer<'a, 'src> {
         self
     }
 
-    fn quoting(&self, string: StringPart) -> Quoting {
+    fn quoting(&self, string: AnyStringPart) -> Quoting {
         match (self.quoting, self.context.f_string_state()) {
             (Quoting::Preserve, _) => Quoting::Preserve,
 
@@ -70,24 +71,21 @@ impl<'a, 'src> StringNormalizer<'a, 'src> {
         }
     }
 
-    /// Computes the strings preferred quotes.
-    pub(crate) fn choose_quotes(&self, string: StringPart) -> QuoteSelection {
-        let raw_content = self.context.locator().slice(string.content_range());
-        let first_quote_or_normalized_char_offset = raw_content
-            .bytes()
-            .position(|b| matches!(b, b'\\' | b'"' | b'\'' | b'\r' | b'{'));
-        let string_flags = string.flags();
-
-        let new_kind = match self.quoting(string) {
-            Quoting::Preserve => string_flags,
+    /// Determines the preferred quote style for `string`.
+    /// The formatter should use the preferred quote style unless
+    /// it can't because the string contains the preferred quotes OR
+    /// it leads to more escaping.
+    pub(super) fn preferred_quote_style(&self, string: AnyStringPart) -> QuoteStyle {
+        match self.quoting(string) {
+            Quoting::Preserve => QuoteStyle::Preserve,
             Quoting::CanChange => {
                 let preferred_quote_style = self
                     .preferred_quote_style
                     .unwrap_or(self.context.options().quote_style());
 
                 // Per PEP 8, always prefer double quotes for triple-quoted strings.
                 // Except when using quote-style-preserve.
-                let preferred_style = if string_flags.is_triple_quoted() {
+                if string.flags().is_triple_quoted() {
                     // ... unless we're formatting a code snippet inside a docstring,
                     // then we specifically want to invert our quote style to avoid
                     // writing out invalid Python.
@@ -142,26 +140,41 @@ impl<'a, 'src> StringNormalizer<'a, 'src> {
                     }
                 } else {
                     preferred_quote_style
-                };
-
-                if let Ok(preferred_quote) = Quote::try_from(preferred_style) {
-                    if let Some(first_quote_or_normalized_char_offset) =
-                        first_quote_or_normalized_char_offset
-                    {
-                        let quote = QuoteMetadata::from_str(
-                            &raw_content[first_quote_or_normalized_char_offset..],
-                            string.flags(),
-                            preferred_quote,
-                        )
-                        .choose(preferred_quote);
-                        string_flags.with_quote_style(quote)
-                    } else {
-                        string_flags.with_quote_style(preferred_quote)
-                    }
-                } else {
-                    string_flags
                 }
             }
+        }
+    }
+
+    /// Computes the strings preferred quotes.
+    pub(crate) fn choose_quotes(&self, string: AnyStringPart) -> QuoteSelection {
+        let raw_content = self.context.locator().slice(string.content_range());
+        let first_quote_or_normalized_char_offset = raw_content
+            .bytes()
+            .position(|b| matches!(b, b'\\' | b'"' | b'\'' | b'\r' | b'{'));
+        let string_flags = string.flags();
+        let preferred_style = self.preferred_quote_style(string);
+
+        let new_kind = match (
+            Quote::try_from(preferred_style),
+            first_quote_or_normalized_char_offset,
+        ) {
+            // The string contains no quotes so it's safe to use the preferred quote style
+            (Ok(preferred_quote), None) => string_flags.with_quote_style(preferred_quote),
+
+            // The preferred quote style is single or double quotes, and the string contains a quote or
+            // another character that may require escaping
+            (Ok(preferred_quote), Some(first_quote_or_normalized_char_offset)) => {
+                let quote = QuoteMetadata::from_str(
+                    &raw_content[first_quote_or_normalized_char_offset..],
+                    string.flags(),
+                    preferred_quote,
+                )
+                .choose(preferred_quote);
+                string_flags.with_quote_style(quote)
+            }
+
+            // The preferred quote style is to preserve the quotes, so let's do that.
+            (Err(_), _) => string_flags,
         };
 
         QuoteSelection {
@@ -171,7 +184,7 @@ impl<'a, 'src> StringNormalizer<'a, 'src> {
     }
 
     /// Computes the strings preferred quotes and normalizes its content.
-    pub(crate) fn normalize(&self, string: StringPart) -> NormalizedString<'src> {
+    pub(crate) fn normalize(&self, string: AnyStringPart) -> NormalizedString<'src> {
         let raw_content = self.context.locator().slice(string.content_range());
         let quote_selection = self.choose_quotes(string);
 
@@ -224,7 +237,7 @@ pub(crate) struct QuoteMetadata {
 /// to choose the quotes for a part.
 impl QuoteMetadata {
     pub(crate) fn from_part(
-        part: StringPart,
+        part: AnyStringPart,
         preferred_quote: Quote,
         context: &PyFormatContext,
     ) -> Self {
@@ -248,7 +261,7 @@ impl QuoteMetadata {
         }
     }
 
-    fn choose(&self, preferred_quote: Quote) -> Quote {
+    pub(super) fn choose(&self, preferred_quote: Quote) -> Quote {
         match self.kind {
             QuoteMetadataKind::Raw { contains_preferred } => {
                 if contains_preferred {
@@ -275,7 +288,7 @@ impl QuoteMetadata {
         }
     }
 
-    fn merge(self, other: &QuoteMetadata) -> Option<QuoteMetadata> {
+    pub(super) fn merge(self, other: &QuoteMetadata) -> Option<QuoteMetadata> {
         match (self.kind, other.kind) {
             (
                 QuoteMetadataKind::Regular {
