Add convenience methods for iterating over all parameter nodes in a function (#11174)

crates/ruff_python_ast/src/helpers.rs

@@ -348,39 +348,18 @@ pub fn any_over_stmt(stmt: &Stmt, func: &dyn Fn(&Expr) -> bool) -> bool {
             returns,
             ..
         }) => {
-            parameters
-                .posonlyargs
-                .iter()
-                .chain(parameters.args.iter().chain(parameters.kwonlyargs.iter()))
-                .any(|parameter| {
-                    parameter
-                        .default
-                        .as_ref()
-                        .is_some_and(|expr| any_over_expr(expr, func))
-                        || parameter
-                            .parameter
-                            .annotation
-                            .as_ref()
-                            .is_some_and(|expr| any_over_expr(expr, func))
-                })
-                || parameters.vararg.as_ref().is_some_and(|parameter| {
-                    parameter
-                        .annotation
-                        .as_ref()
-                        .is_some_and(|expr| any_over_expr(expr, func))
-                })
-                || parameters.kwarg.as_ref().is_some_and(|parameter| {
-                    parameter
-                        .annotation
-                        .as_ref()
-                        .is_some_and(|expr| any_over_expr(expr, func))
-                })
-                || type_params.as_ref().is_some_and(|type_params| {
-                    type_params
-                        .iter()
-                        .any(|type_param| any_over_type_param(type_param, func))
-                })
-                || body.iter().any(|stmt| any_over_stmt(stmt, func))
+            parameters.iter().any(|param| {
+                param
+                    .default()
+                    .is_some_and(|default| any_over_expr(default, func))
+                    || param
+                        .annotation()
+                        .is_some_and(|annotation| any_over_expr(annotation, func))
+            }) || type_params.as_ref().is_some_and(|type_params| {
+                type_params
+                    .iter()
+                    .any(|type_param| any_over_type_param(type_param, func))
+            }) || body.iter().any(|stmt| any_over_stmt(stmt, func))
                 || decorator_list
                     .iter()
                     .any(|decorator| any_over_expr(&decorator.expression, func))

crates/ruff_python_ast/src/node.rs

@@ -1,12 +1,13 @@
 use crate::visitor::preorder::PreorderVisitor;
 use crate::{
-    self as ast, Alias, ArgOrKeyword, Arguments, Comprehension, Decorator, ExceptHandler, Expr,
-    FStringElement, Keyword, MatchCase, Mod, Parameter, ParameterWithDefault, Parameters, Pattern,
-    PatternArguments, PatternKeyword, Stmt, StmtAnnAssign, StmtAssert, StmtAssign, StmtAugAssign,
-    StmtBreak, StmtClassDef, StmtContinue, StmtDelete, StmtExpr, StmtFor, StmtFunctionDef,
-    StmtGlobal, StmtIf, StmtImport, StmtImportFrom, StmtIpyEscapeCommand, StmtMatch, StmtNonlocal,
-    StmtPass, StmtRaise, StmtReturn, StmtTry, StmtTypeAlias, StmtWhile, StmtWith, TypeParam,
-    TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams, WithItem,
+    self as ast, Alias, AnyParameterRef, ArgOrKeyword, Arguments, Comprehension, Decorator,
+    ExceptHandler, Expr, FStringElement, Keyword, MatchCase, Mod, Parameter, ParameterWithDefault,
+    Parameters, Pattern, PatternArguments, PatternKeyword, Stmt, StmtAnnAssign, StmtAssert,
+    StmtAssign, StmtAugAssign, StmtBreak, StmtClassDef, StmtContinue, StmtDelete, StmtExpr,
+    StmtFor, StmtFunctionDef, StmtGlobal, StmtIf, StmtImport, StmtImportFrom, StmtIpyEscapeCommand,
+    StmtMatch, StmtNonlocal, StmtPass, StmtRaise, StmtReturn, StmtTry, StmtTypeAlias, StmtWhile,
+    StmtWith, TypeParam, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams,
+    WithItem,
 };
 use ruff_text_size::{Ranged, TextRange};
 use std::ptr::NonNull;
@@ -4221,28 +4222,13 @@ impl AstNode for Parameters {
     where
         V: PreorderVisitor<'a> + ?Sized,
     {
-        let ast::Parameters {
-            range: _,
-            posonlyargs,
-            args,
-            vararg,
-            kwonlyargs,
-            kwarg,
-        } = self;
-        for arg in posonlyargs.iter().chain(args) {
-            visitor.visit_parameter_with_default(arg);
-        }
-
-        if let Some(arg) = vararg {
-            visitor.visit_parameter(arg);
-        }
-
-        for arg in kwonlyargs {
-            visitor.visit_parameter_with_default(arg);
-        }
-
-        if let Some(arg) = kwarg {
-            visitor.visit_parameter(arg);
+        for parameter in self {
+            match parameter {
+                AnyParameterRef::NonVariadic(parameter_with_default) => {
+                    visitor.visit_parameter_with_default(parameter_with_default);
+                }
+                AnyParameterRef::Variadic(parameter) => visitor.visit_parameter(parameter),
+            }
         }
     }
 }

crates/ruff_python_ast/src/nodes.rs

@@ -2,6 +2,7 @@
 
 use std::fmt;
 use std::fmt::Debug;
+use std::iter::FusedIterator;
 use std::ops::Deref;
 use std::slice::{Iter, IterMut};
 use std::sync::OnceLock;
@@ -3175,6 +3176,63 @@ pub struct Decorator {
     pub expression: Expr,
 }
 
+/// Enumeration of the two kinds of parameter
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum AnyParameterRef<'a> {
+    /// Variadic parameters cannot have default values,
+    /// e.g. both `*args` and `**kwargs` in the following function:
+    ///
+    /// ```python
+    /// def foo(*args, **kwargs): pass
+    /// ```
+    Variadic(&'a Parameter),
+
+    /// Non-variadic parameters can have default values,
+    /// though they won't necessarily always have them:
+    ///
+    /// ```python
+    /// def bar(a=1, /, b=2, *, c=3): pass
+    /// ```
+    NonVariadic(&'a ParameterWithDefault),
+}
+
+impl<'a> AnyParameterRef<'a> {
+    pub const fn as_parameter(self) -> &'a Parameter {
+        match self {
+            Self::NonVariadic(param) => &param.parameter,
+            Self::Variadic(param) => param,
+        }
+    }
+
+    pub const fn name(self) -> &'a Identifier {
+        &self.as_parameter().name
+    }
+
+    pub const fn is_variadic(self) -> bool {
+        matches!(self, Self::Variadic(_))
+    }
+
+    pub fn annotation(self) -> Option<&'a Expr> {
+        self.as_parameter().annotation.as_deref()
+    }
+
+    pub fn default(self) -> Option<&'a Expr> {
+        match self {
+            Self::NonVariadic(param) => param.default.as_deref(),
+            Self::Variadic(_) => None,
+        }
+    }
+}
+
+impl Ranged for AnyParameterRef<'_> {
+    fn range(&self) -> TextRange {
+        match self {
+            Self::NonVariadic(param) => param.range,
+            Self::Variadic(param) => param.range,
+        }
+    }
+}
+
 /// An alternative type of AST `arguments`. This is ruff_python_parser-friendly and human-friendly definition of function arguments.
 /// This form also has advantage to implement pre-order traverse.
 ///
@@ -3196,37 +3254,56 @@ pub struct Parameters {
 }
 
 impl Parameters {
-    /// Returns the [`ParameterWithDefault`] with the given name, or `None` if no such [`ParameterWithDefault`] exists.
-    pub fn find(&self, name: &str) -> Option<&ParameterWithDefault> {
+    /// Returns an iterator over all non-variadic parameters included in this [`Parameters`] node.
+    ///
+    /// The variadic parameters (`.vararg` and `.kwarg`) can never have default values;
+    /// non-variadic parameters sometimes will.
+    pub fn iter_non_variadic_params(&self) -> impl Iterator<Item = &ParameterWithDefault> {
         self.posonlyargs
             .iter()
             .chain(&self.args)
             .chain(&self.kwonlyargs)
+    }
+
+    /// Returns the [`ParameterWithDefault`] with the given name, or `None` if no such [`ParameterWithDefault`] exists.
+    pub fn find(&self, name: &str) -> Option<&ParameterWithDefault> {
+        self.iter_non_variadic_params()
             .find(|arg| arg.parameter.name.as_str() == name)
     }
 
-    /// Returns `true` if a parameter with the given name included in this [`Parameters`].
+    /// Returns an iterator over all parameters included in this [`Parameters`] node.
+    pub fn iter(&self) -> ParametersIterator {
+        ParametersIterator::new(self)
+    }
+
+    /// Returns the total number of parameters included in this [`Parameters`] node.
+    pub fn len(&self) -> usize {
+        let Parameters {
+            range: _,
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kwarg,
+        } = self;
+        // Safety: a Python function can have an arbitrary number of parameters,
+        // so theoretically this could be a number that wouldn't fit into a usize,
+        // which would lead to a panic. A Python function with that many parameters
+        // is extremely unlikely outside of generated code, however, and it's even
+        // more unlikely that we'd find a function with that many parameters in a
+        // source-code file <=4GB large (Ruff's maximum).
+        posonlyargs
+            .len()
+            .checked_add(args.len())
+            .and_then(|length| length.checked_add(usize::from(vararg.is_some())))
+            .and_then(|length| length.checked_add(kwonlyargs.len()))
+            .and_then(|length| length.checked_add(usize::from(kwarg.is_some())))
+            .expect("Failed to fit the number of parameters into a usize")
+    }
+
+    /// Returns `true` if a parameter with the given name is included in this [`Parameters`].
     pub fn includes(&self, name: &str) -> bool {
-        if self
-            .posonlyargs
-            .iter()
-            .chain(&self.args)
-            .chain(&self.kwonlyargs)
-            .any(|arg| arg.parameter.name.as_str() == name)
-        {
-            return true;
-        }
-        if let Some(arg) = &self.vararg {
-            if arg.name.as_str() == name {
-                return true;
-            }
-        }
-        if let Some(arg) = &self.kwarg {
-            if arg.name.as_str() == name {
-                return true;
-            }
-        }
-        false
+        self.iter().any(|param| param.name() == name)
     }
 
     /// Returns `true` if the [`Parameters`] is empty.
@@ -3239,6 +3316,136 @@ impl Parameters {
     }
 }
 
+pub struct ParametersIterator<'a> {
+    posonlyargs: Iter<'a, ParameterWithDefault>,
+    args: Iter<'a, ParameterWithDefault>,
+    vararg: Option<&'a Parameter>,
+    kwonlyargs: Iter<'a, ParameterWithDefault>,
+    kwarg: Option<&'a Parameter>,
+}
+
+impl<'a> ParametersIterator<'a> {
+    fn new(parameters: &'a Parameters) -> Self {
+        let Parameters {
+            range: _,
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kwarg,
+        } = parameters;
+        Self {
+            posonlyargs: posonlyargs.iter(),
+            args: args.iter(),
+            vararg: vararg.as_deref(),
+            kwonlyargs: kwonlyargs.iter(),
+            kwarg: kwarg.as_deref(),
+        }
+    }
+}
+
+impl<'a> Iterator for ParametersIterator<'a> {
+    type Item = AnyParameterRef<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let ParametersIterator {
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kwarg,
+        } = self;
+
+        if let Some(param) = posonlyargs.next() {
+            return Some(AnyParameterRef::NonVariadic(param));
+        }
+        if let Some(param) = args.next() {
+            return Some(AnyParameterRef::NonVariadic(param));
+        }
+        if let Some(param) = vararg.take() {
+            return Some(AnyParameterRef::Variadic(param));
+        }
+        if let Some(param) = kwonlyargs.next() {
+            return Some(AnyParameterRef::NonVariadic(param));
+        }
+        kwarg.take().map(AnyParameterRef::Variadic)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let ParametersIterator {
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kwarg,
+        } = self;
+
+        let posonlyargs_len = posonlyargs.len();
+        let args_len = args.len();
+        let vararg_len = usize::from(vararg.is_some());
+        let kwonlyargs_len = kwonlyargs.len();
+        let kwarg_len = usize::from(kwarg.is_some());
+
+        let lower = posonlyargs_len
+            .saturating_add(args_len)
+            .saturating_add(vararg_len)
+            .saturating_add(kwonlyargs_len)
+            .saturating_add(kwarg_len);
+
+        let upper = posonlyargs_len
+            .checked_add(args_len)
+            .and_then(|length| length.checked_add(vararg_len))
+            .and_then(|length| length.checked_add(kwonlyargs_len))
+            .and_then(|length| length.checked_add(kwarg_len));
+
+        (lower, upper)
+    }
+
+    fn last(mut self) -> Option<Self::Item> {
+        self.next_back()
+    }
+}
+
+impl<'a> DoubleEndedIterator for ParametersIterator<'a> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        let ParametersIterator {
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kwarg,
+        } = self;
+
+        if let Some(param) = kwarg.take() {
+            return Some(AnyParameterRef::Variadic(param));
+        }
+        if let Some(param) = kwonlyargs.next_back() {
+            return Some(AnyParameterRef::NonVariadic(param));
+        }
+        if let Some(param) = vararg.take() {
+            return Some(AnyParameterRef::Variadic(param));
+        }
+        if let Some(param) = args.next_back() {
+            return Some(AnyParameterRef::NonVariadic(param));
+        }
+        posonlyargs.next_back().map(AnyParameterRef::NonVariadic)
+    }
+}
+
+impl<'a> FusedIterator for ParametersIterator<'a> {}
+
+/// We rely on the same invariants outlined in the comment above `Parameters::len()`
+/// in order to implement `ExactSizeIterator` here
+impl<'a> ExactSizeIterator for ParametersIterator<'a> {}
+
+impl<'a> IntoIterator for &'a Parameters {
+    type IntoIter = ParametersIterator<'a>;
+    type Item = AnyParameterRef<'a>;
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
 /// An alternative type of AST `arg`. This is used for each function argument that might have a default value.
 /// Used by `Arguments` original type.
 ///

crates/ruff_python_ast/src/visitor.rs

@@ -4,11 +4,11 @@ pub mod preorder;
 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, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple,
-    TypeParams, UnaryOp, WithItem,
+    self as ast, Alias, AnyParameterRef, Arguments, BoolOp, BytesLiteral, CmpOp, Comprehension,
+    Decorator, ElifElseClause, ExceptHandler, Expr, ExprContext, FString, FStringElement,
+    FStringPart, Keyword, MatchCase, Operator, Parameter, Parameters, Pattern, PatternArguments,
+    PatternKeyword, Stmt, StringLiteral, TypeParam, TypeParamParamSpec, TypeParamTypeVar,
+    TypeParamTypeVarTuple, TypeParams, UnaryOp, WithItem,
 };
 
 /// A trait for AST visitors. Visits all nodes in the AST recursively in evaluation-order.
@@ -607,36 +607,15 @@ pub fn walk_arguments<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, arguments: &
 
 pub fn walk_parameters<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, parameters: &'a Parameters) {
     // Defaults are evaluated before annotations.
-    for arg in &parameters.posonlyargs {
-        if let Some(default) = &arg.default {
-            visitor.visit_expr(default);
-        }
-    }
-    for arg in &parameters.args {
-        if let Some(default) = &arg.default {
-            visitor.visit_expr(default);
-        }
-    }
-    for arg in &parameters.kwonlyargs {
-        if let Some(default) = &arg.default {
-            visitor.visit_expr(default);
-        }
+    for default in parameters
+        .iter_non_variadic_params()
+        .filter_map(|param| param.default.as_deref())
+    {
+        visitor.visit_expr(default);
     }
 
-    for arg in &parameters.posonlyargs {
-        visitor.visit_parameter(&arg.parameter);
-    }
-    for arg in &parameters.args {
-        visitor.visit_parameter(&arg.parameter);
-    }
-    if let Some(arg) = &parameters.vararg {
-        visitor.visit_parameter(arg);
-    }
-    for arg in &parameters.kwonlyargs {
-        visitor.visit_parameter(&arg.parameter);
-    }
-    if let Some(arg) = &parameters.kwarg {
-        visitor.visit_parameter(arg);
+    for parameter in parameters.iter().map(AnyParameterRef::as_parameter) {
+        visitor.visit_parameter(parameter);
     }
 }