Update Stylist quote detection with new f-string token

crates/ruff_python_codegen/src/stylist.rs

@@ -9,7 +9,7 @@ use ruff_python_parser::lexer::LexResult;
 use ruff_python_parser::Tok;
 use ruff_source_file::{find_newline, LineEnding};
 
-use ruff_python_ast::str::leading_quote;
+use ruff_python_ast::str::{is_triple_quote, leading_quote};
 use ruff_source_file::Locator;
 
 pub struct Stylist<'a> {
@@ -55,6 +55,13 @@ fn detect_quote(tokens: &[LexResult], locator: &Locator) -> Quote {
  triple_quoted: false,
  ..
  } => Some(*range),
+ Tok::FStringStart => {
+ if is_triple_quote(locator.slice(range)) {
+ None
+ } else {
+ Some(*range)
+ }
+ }
  _ => None,
  });
 
@@ -248,6 +255,14 @@ x = (
  Quote::Single
  );
 
+ let contents = r#"x = f'1'"#;
+ let locator = Locator::new(contents);
+ let tokens: Vec<_> = lex(contents, Mode::Module).collect();
+ assert_eq!(
+ Stylist::from_tokens(&tokens, &locator).quote(),
+ Quote::Single
+ );
+
  let contents = r#"x = "1""#;
  let locator = Locator::new(contents);
  let tokens: Vec<_> = lex(contents, Mode::Module).collect();
@@ -256,6 +271,14 @@ x = (
  Quote::Double
  );
 
+ let contents = r#"x = f"1""#;
+ let locator = Locator::new(contents);
+ let tokens: Vec<_> = lex(contents, Mode::Module).collect();
+ assert_eq!(
+ Stylist::from_tokens(&tokens, &locator).quote(),
+ Quote::Double
+ );
+
  let contents = r#"s = "It's done.""#;
  let locator = Locator::new(contents);
  let tokens: Vec<_> = lex(contents, Mode::Module).collect();
@@ -294,6 +317,31 @@ a = 'v'
 '''Module docstring.'''
 
 a = "v"
+"#;
+ let locator = Locator::new(contents);
+ let tokens: Vec<_> = lex(contents, Mode::Module).collect();
+ assert_eq!(
+ Stylist::from_tokens(&tokens, &locator).quote(),
+ Quote::Double
+ );
+
+ // Detect from f-string appearing after docstring
+ let contents = r#"
+"""Module docstring."""
+
+a = f'v'
+"#;
+ let locator = Locator::new(contents);
+ let tokens: Vec<_> = lex(contents, Mode::Module).collect();
+ assert_eq!(
+ Stylist::from_tokens(&tokens, &locator).quote(),
+ Quote::Single
+ );
+
+ let contents = r#"
+'''Module docstring.'''
+
+a = f"v"
 "#;
  let locator = Locator::new(contents);
  let tokens: Vec<_> = lex(contents, Mode::Module).collect();