Invert quote-style when generating code within f-strings (#4487)

crates/ruff_python_ast/src/helpers.rs

@@ -17,26 +17,23 @@ use smallvec::SmallVec;
 
 use crate::call_path::CallPath;
 use crate::newlines::UniversalNewlineIterator;
-use crate::source_code::{Generator, Indexer, Locator, Stylist};
+use crate::source_code::{Generator, Indexer, Locator};
 use crate::statement_visitor::{walk_body, walk_stmt, StatementVisitor};
 
 /// Generate source code from an [`Expr`].
-pub fn unparse_expr(expr: &Expr, stylist: &Stylist) -> String {
-    let mut generator: Generator = stylist.into();
+pub fn unparse_expr(expr: &Expr, mut generator: Generator) -> String {
     generator.unparse_expr(expr, 0);
     generator.generate()
 }
 
 /// Generate source code from a [`Stmt`].
-pub fn unparse_stmt(stmt: &Stmt, stylist: &Stylist) -> String {
-    let mut generator: Generator = stylist.into();
+pub fn unparse_stmt(stmt: &Stmt, mut generator: Generator) -> String {
     generator.unparse_stmt(stmt);
     generator.generate()
 }
 
 /// Generate source code from an [`Constant`].
-pub fn unparse_constant(constant: &Constant, stylist: &Stylist) -> String {
-    let mut generator: Generator = stylist.into();
+pub fn unparse_constant(constant: &Constant, mut generator: Generator) -> String {
     generator.unparse_constant(constant);
     generator.generate()
 }

crates/ruff_python_ast/src/source_code/indexer.rs

@@ -4,7 +4,7 @@
 use crate::source_code::Locator;
 use ruff_text_size::{TextRange, TextSize};
 use rustpython_parser::lexer::LexResult;
-use rustpython_parser::Tok;
+use rustpython_parser::{StringKind, Tok};
 
 pub struct Indexer {
     /// Stores the ranges of comments sorted by [`TextRange::start`] in increasing order. No two ranges are overlapping.
@@ -16,15 +16,20 @@ pub struct Indexer {
     /// The range of all triple quoted strings in the source document. The ranges are sorted by their
     /// [`TextRange::start`] position in increasing order. No two ranges are overlapping.
     triple_quoted_string_ranges: Vec<TextRange>,
+
+    /// The range of all f-string in the source document. The ranges are sorted by their
+    /// [`TextRange::start`] position in increasing order. No two ranges are overlapping.
+    f_string_ranges: Vec<TextRange>,
 }
 
 impl Indexer {
     pub fn from_tokens(tokens: &[LexResult], locator: &Locator) -> Self {
         assert!(TextSize::try_from(locator.contents().len()).is_ok());
 
-        let mut commented_lines = Vec::new();
+        let mut comment_ranges = Vec::new();
         let mut continuation_lines = Vec::new();
-        let mut string_ranges = Vec::new();
+        let mut triple_quoted_string_ranges = Vec::new();
+        let mut f_string_ranges = Vec::new();
         // Token, end
         let mut prev_end = TextSize::default();
         let mut prev_token: Option<&Tok> = None;
@@ -59,7 +64,7 @@ impl Indexer {
 
             match tok {
                 Tok::Comment(..) => {
-                    commented_lines.push(*range);
+                    comment_ranges.push(*range);
                 }
                 Tok::Newline | Tok::NonLogicalNewline => {
                     line_start = range.end();
@@ -67,7 +72,15 @@ impl Indexer {
                 Tok::String {
                     triple_quoted: true,
                     ..
-                } => string_ranges.push(*range),
+                } => {
+                    triple_quoted_string_ranges.push(*range);
+                }
+                Tok::String {
+                    kind: StringKind::FString | StringKind::RawFString,
+                    ..
+                } => {
+                    f_string_ranges.push(*range);
+                }
                 _ => {}
             }
 
@@ -75,9 +88,10 @@ impl Indexer {
             prev_end = range.end();
         }
         Self {
-            comment_ranges: commented_lines,
+            comment_ranges,
             continuation_lines,
-            triple_quoted_string_ranges: string_ranges,
+            triple_quoted_string_ranges,
+            f_string_ranges,
         }
     }
 
@@ -97,10 +111,27 @@ impl Indexer {
         &self.triple_quoted_string_ranges
     }
 
+    /// Returns `true` if the given offset is part of a continuation line.
     pub fn is_continuation(&self, offset: TextSize, locator: &Locator) -> bool {
         let line_start = locator.line_start(offset);
         self.continuation_lines.binary_search(&line_start).is_ok()
     }
+
+    /// Return the [`TextRange`] of the f-string containing a given offset.
+    pub fn f_string_range(&self, offset: TextSize) -> Option<TextRange> {
+        let Ok(string_range_index) = self.f_string_ranges.binary_search_by(|range| {
+            if offset < range.start() {
+                std::cmp::Ordering::Greater
+            } else if range.contains(offset) {
+                std::cmp::Ordering::Equal
+            } else {
+                std::cmp::Ordering::Less
+            }
+        }) else {
+            return None;
+        };
+        Some(self.f_string_ranges[string_range_index])
+    }
 }
 
 #[cfg(test)]

crates/ruff_python_ast/src/source_code/mod.rs

@@ -15,7 +15,7 @@ use rustpython_parser::{lexer, Mode, ParseError};
 use serde::{Deserialize, Serialize};
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
-pub use stylist::Stylist;
+pub use stylist::{Quote, Stylist};
 
 /// Run round-trip source code generation on a given Python code.
 pub fn round_trip(code: &str, source_path: &str) -> Result<String, ParseError> {