From 5c23131beaf71151ae1dc60154dde8986208bd92 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 13 Sep 2023 10:26:39 +0530
Subject: [PATCH] Update `Stylist` quote detection with new f-string token

---
 crates/ruff_python_codegen/src/stylist.rs | 50 ++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/crates/ruff_python_codegen/src/stylist.rs b/crates/ruff_python_codegen/src/stylist.rs
index 087d82a29abe76..054da50fb3fed5 100644
--- a/crates/ruff_python_codegen/src/stylist.rs
+++ b/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();