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

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.
 ///