From 658435ed9eeaa4458562235029265c57f858ed2d Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 15 Sep 2023 12:52:26 +0530
Subject: [PATCH] Update `Stylist` quote detection with new f-string token
 (#7328)

## Summary

This PR updates `Stylist` quote detection to include the f-string
tokens.

As f-strings cannot be used as docstrings, we'll skip the check for
triple-quoted f-strings.

## Test Plan

Add new test cases with f-strings.

fixes: #7293
---
 crates/ruff_python_codegen/src/stylist.rs | 54 +++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/crates/ruff_python_codegen/src/stylist.rs b/crates/ruff_python_codegen/src/stylist.rs
index 4e4f92227b812..1c3740bbcef76 100644
--- a/crates/ruff_python_codegen/src/stylist.rs
+++ b/crates/ruff_python_codegen/src/stylist.rs
@@ -55,6 +55,9 @@ fn detect_quote(tokens: &[LexResult], locator: &Locator) -> Quote {
             triple_quoted: false,
             ..
         } => Some(*range),
+        // No need to check if it's triple-quoted as f-strings cannot be used
+        // as docstrings.
+        Tok::FStringStart => Some(*range),
         _ => None,
     });
 
@@ -275,6 +278,14 @@ class FormFeedIndent:
             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();
@@ -283,6 +294,14 @@ class FormFeedIndent:
             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();
@@ -328,6 +347,41 @@ a = "v"
             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();
+        assert_eq!(
+            Stylist::from_tokens(&tokens, &locator).quote(),
+            Quote::Double
+        );
+
+        let contents = r#"
+f'''Module docstring.'''
+"#;
+        let locator = Locator::new(contents);
+        let tokens: Vec<_> = lex(contents, Mode::Module).collect();
+        assert_eq!(
+            Stylist::from_tokens(&tokens, &locator).quote(),
+            Quote::Single
+        );
     }
 
     #[test]