Add visitor method for FStringFormatSpec

crates/ruff_python_ast/src/node.rs

@@ -2712,9 +2712,7 @@ impl AstNode for ast::FStringExpressionElement {
         visitor.visit_expr(expression);
 
         if let Some(format_spec) = format_spec {
-            for spec_part in &format_spec.elements {
-                visitor.visit_f_string_element(spec_part);
-            }
+            visitor.visit_f_string_format_spec(format_spec);
         }
     }
 }

crates/ruff_python_ast/src/visitor.rs

@@ -5,9 +5,10 @@ pub mod transformer;
 
 use crate::{
     self as ast, Alias, Arguments, BoolOp, BytesLiteral, CmpOp, Comprehension, Decorator,
-    ElifElseClause, ExceptHandler, Expr, ExprContext, FString, FStringElement, FStringPart,
-    Keyword, MatchCase, Operator, Parameter, Parameters, Pattern, PatternArguments, PatternKeyword,
-    Stmt, StringLiteral, TypeParam, TypeParamTypeVar, TypeParams, UnaryOp, WithItem,
+    ElifElseClause, ExceptHandler, Expr, ExprContext, FString, FStringElement, FStringFormatSpec,
+    FStringPart, Keyword, MatchCase, Operator, Parameter, Parameters, Pattern, PatternArguments,
+    PatternKeyword, Stmt, StringLiteral, TypeParam, TypeParamTypeVar, TypeParams, UnaryOp,
+    WithItem,
 };
 
 /// A trait for AST visitors. Visits all nodes in the AST recursively in evaluation-order.
@@ -101,6 +102,9 @@ pub trait Visitor<'a> {
     fn visit_f_string_element(&mut self, f_string_element: &'a FStringElement) {
         walk_f_string_element(self, f_string_element);
     }
+    fn visit_f_string_format_spec(&mut self, f_string_format_spec: &'a FStringFormatSpec) {
+        walk_f_string_format_spec(self, f_string_format_spec);
+    }
     fn visit_string_literal(&mut self, string_literal: &'a StringLiteral) {
         walk_string_literal(self, string_literal);
     }
@@ -745,6 +749,15 @@ pub fn walk_f_string<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, f_string: &'a
     }
 }
 
+pub fn walk_f_string_format_spec<'a, V: Visitor<'a> + ?Sized>(
+    visitor: &mut V,
+    f_string_format_spec: &'a FStringFormatSpec,
+) {
+    for f_string_element in &f_string_format_spec.elements {
+        visitor.visit_f_string_element(f_string_element);
+    }
+}
+
 pub fn walk_f_string_element<'a, V: Visitor<'a> + ?Sized>(
     visitor: &mut V,
     f_string_element: &'a FStringElement,
@@ -757,9 +770,7 @@ pub fn walk_f_string_element<'a, V: Visitor<'a> + ?Sized>(
     {
         visitor.visit_expr(expression);
         if let Some(format_spec) = format_spec {
-            for spec_element in &format_spec.elements {
-                visitor.visit_f_string_element(spec_element);
-            }
+            visitor.visit_f_string_format_spec(format_spec);
         }
     }
 }

crates/ruff_python_ast/src/visitor/preorder.rs

@@ -1,8 +1,8 @@
 use crate::{
     Alias, Arguments, BoolOp, BytesLiteral, CmpOp, Comprehension, Decorator, ElifElseClause,
-    ExceptHandler, Expr, FString, FStringElement, Keyword, MatchCase, Mod, Operator, Parameter,
-    ParameterWithDefault, Parameters, Pattern, PatternArguments, PatternKeyword, Singleton, Stmt,
-    StringLiteral, TypeParam, TypeParams, UnaryOp, WithItem,
+    ExceptHandler, Expr, FString, FStringElement, FStringFormatSpec, Keyword, MatchCase, Mod,
+    Operator, Parameter, ParameterWithDefault, Parameters, Pattern, PatternArguments,
+    PatternKeyword, Singleton, Stmt, StringLiteral, TypeParam, TypeParams, UnaryOp, WithItem,
 };
 use crate::{AnyNodeRef, AstNode};
 
@@ -158,6 +158,11 @@ pub trait PreorderVisitor<'a> {
         walk_f_string_element(self, f_string_element);
     }
 
+    #[inline]
+    fn visit_f_string_format_spec(&mut self, f_string_format_spec: &'a FStringFormatSpec) {
+        walk_f_string_format_spec(self, f_string_format_spec);
+    }
+
     #[inline]
     fn visit_string_literal(&mut self, string_literal: &'a StringLiteral) {
         walk_string_literal(self, string_literal);
@@ -570,6 +575,20 @@ where
     visitor.leave_node(node);
 }
 
+#[inline]
+pub fn walk_f_string_format_spec<'a, V>(
+    visitor: &mut V,
+    f_string_format_spec: &'a FStringFormatSpec,
+) where
+    V: PreorderVisitor<'a> + ?Sized,
+{
+    let node = AnyNodeRef::from(f_string_format_spec);
+    if visitor.enter_node(node).is_traverse() {
+        f_string_format_spec.visit_preorder(visitor);
+    }
+    visitor.leave_node(node);
+}
+
 #[inline]
 pub fn walk_string_literal<'a, V>(visitor: &mut V, string_literal: &'a StringLiteral)
 where

crates/ruff_python_ast/src/visitor/transformer.rs

@@ -1,8 +1,8 @@
 use crate::{
     self as ast, Alias, Arguments, BoolOp, BytesLiteral, CmpOp, Comprehension, Decorator,
-    ElifElseClause, ExceptHandler, Expr, ExprContext, FString, FStringElement, Keyword, MatchCase,
-    Operator, Parameter, Parameters, Pattern, PatternArguments, PatternKeyword, Stmt,
-    StringLiteral, TypeParam, TypeParamTypeVar, TypeParams, UnaryOp, WithItem,
+    ElifElseClause, ExceptHandler, Expr, ExprContext, FString, FStringElement, FStringFormatSpec,
+    Keyword, MatchCase, Operator, Parameter, Parameters, Pattern, PatternArguments, PatternKeyword,
+    Stmt, StringLiteral, TypeParam, TypeParamTypeVar, TypeParams, UnaryOp, WithItem,
 };
 
 /// A trait for transforming ASTs. Visits all nodes in the AST recursively in evaluation-order.
@@ -88,6 +88,9 @@ pub trait Transformer {
     fn visit_f_string_element(&self, f_string_element: &mut FStringElement) {
         walk_f_string_element(self, f_string_element);
     }
+    fn visit_f_string_format_spec(&self, f_string_format_spec: &mut FStringFormatSpec) {
+        walk_f_string_format_spec(self, f_string_format_spec);
+    }
     fn visit_string_literal(&self, string_literal: &mut StringLiteral) {
         walk_string_literal(self, string_literal);
     }
@@ -731,6 +734,15 @@ pub fn walk_f_string<V: Transformer + ?Sized>(visitor: &V, f_string: &mut FStrin
     }
 }
 
+pub fn walk_f_string_format_spec<V: Transformer + ?Sized>(
+    visitor: &V,
+    f_string_format_spec: &mut FStringFormatSpec,
+) {
+    for spec_element in &mut f_string_format_spec.elements {
+        visitor.visit_f_string_element(spec_element);
+    }
+}
+
 pub fn walk_f_string_element<V: Transformer + ?Sized>(
     visitor: &V,
     f_string_element: &mut FStringElement,
@@ -743,9 +755,7 @@ pub fn walk_f_string_element<V: Transformer + ?Sized>(
     {
         visitor.visit_expr(expression);
         if let Some(format_spec) = format_spec {
-            for spec_element in &mut format_spec.elements {
-                visitor.visit_f_string_element(spec_element);
-            }
+            visitor.visit_f_string_format_spec(format_spec);
         }
     }
 }

crates/ruff_python_ast/tests/snapshots/preorder__f_strings.snap

@@ -10,9 +10,10 @@ expression: trace
         - FStringLiteralElement
         - FStringExpressionElement
           - ExprName
-          - FStringLiteralElement
-          - FStringExpressionElement
-            - ExprName
-          - FStringLiteralElement
+          - FStringFormatSpec
+            - FStringLiteralElement
+            - FStringExpressionElement
+              - ExprName
+            - FStringLiteralElement
         - FStringLiteralElement