From 32d6f84e3d6b52a160db0af0b5819848e74f11ea Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 13 Mar 2024 14:16:55 +0530
Subject: [PATCH] Add methods to iter over f-string elements (#10309)

## Summary

This PR adds methods on `FString` to iterate over the two different kind
of elements it can have - literals and expressions. This is similar to
the methods we have on `ExprFString`.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 .../rules/ruff/rules/missing_fstring_syntax.rs   |  6 +-----
 .../tryceratops/rules/raise_vanilla_args.rs      |  6 +-----
 crates/ruff_python_ast/src/nodes.rs              | 16 ++++++++++++++++
 .../ruff_python_formatter/src/other/f_string.rs  |  4 +---
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/crates/ruff_linter/src/rules/ruff/rules/missing_fstring_syntax.rs b/crates/ruff_linter/src/rules/ruff/rules/missing_fstring_syntax.rs
index 583a51947398d..09cd172cf715c 100644
--- a/crates/ruff_linter/src/rules/ruff/rules/missing_fstring_syntax.rs
+++ b/crates/ruff_linter/src/rules/ruff/rules/missing_fstring_syntax.rs
@@ -168,11 +168,7 @@ fn should_be_fstring(
 
     for f_string in value.f_strings() {
         let mut has_name = false;
-        for element in f_string
-            .elements
-            .iter()
-            .filter_map(|element| element.as_expression())
-        {
+        for element in f_string.expressions() {
             if let ast::Expr::Name(ast::ExprName { id, .. }) = element.expression.as_ref() {
                 if arg_names.contains(id.as_str()) {
                     return false;
diff --git a/crates/ruff_linter/src/rules/tryceratops/rules/raise_vanilla_args.rs b/crates/ruff_linter/src/rules/tryceratops/rules/raise_vanilla_args.rs
index 561f2f9577599..92d09b15a97d8 100644
--- a/crates/ruff_linter/src/rules/tryceratops/rules/raise_vanilla_args.rs
+++ b/crates/ruff_linter/src/rules/tryceratops/rules/raise_vanilla_args.rs
@@ -100,11 +100,7 @@ fn contains_message(expr: &Expr) -> bool {
                         }
                     }
                     ast::FStringPart::FString(f_string) => {
-                        for literal in f_string
-                            .elements
-                            .iter()
-                            .filter_map(|element| element.as_literal())
-                        {
+                        for literal in f_string.literals() {
                             if literal.chars().any(char::is_whitespace) {
                                 return true;
                             }
diff --git a/crates/ruff_python_ast/src/nodes.rs b/crates/ruff_python_ast/src/nodes.rs
index 53cf5b8d809a6..89133331f1719 100644
--- a/crates/ruff_python_ast/src/nodes.rs
+++ b/crates/ruff_python_ast/src/nodes.rs
@@ -1248,6 +1248,22 @@ pub struct FString {
     pub flags: FStringFlags,
 }
 
+impl FString {
+    /// Returns an iterator over all the [`FStringLiteralElement`] nodes contained in this f-string.
+    pub fn literals(&self) -> impl Iterator<Item = &FStringLiteralElement> {
+        self.elements
+            .iter()
+            .filter_map(|element| element.as_literal())
+    }
+
+    /// Returns an iterator over all the [`FStringExpressionElement`] nodes contained in this f-string.
+    pub fn expressions(&self) -> impl Iterator<Item = &FStringExpressionElement> {
+        self.elements
+            .iter()
+            .filter_map(|element| element.as_expression())
+    }
+}
+
 impl Ranged for FString {
     fn range(&self) -> TextRange {
         self.range
diff --git a/crates/ruff_python_formatter/src/other/f_string.rs b/crates/ruff_python_formatter/src/other/f_string.rs
index aa78f7520bf81..0bae84a1d1832 100644
--- a/crates/ruff_python_formatter/src/other/f_string.rs
+++ b/crates/ruff_python_formatter/src/other/f_string.rs
@@ -138,9 +138,7 @@ impl FStringLayout {
         //
         // Reference: https://prettier.io/docs/en/next/rationale.html#template-literals
         if f_string
-            .elements
-            .iter()
-            .filter_map(|element| element.as_expression())
+            .expressions()
             .any(|expr| memchr::memchr2(b'\n', b'\r', locator.slice(expr).as_bytes()).is_some())
         {
             Self::Multiline