Add a helper for any-like operations (#1757)

src/ast/helpers.rs

@@ -187,86 +187,61 @@ pub fn contains_call_path(
     import_aliases: &FxHashMap<&str, &str>,
     from_imports: &FxHashMap<&str, FxHashSet<&str>>,
 ) -> bool {
-    let call_path = collect_call_paths(expr);
+    any_over_expr(expr, &|expr| {
-    if !call_path.is_empty() {
+        let call_path = collect_call_paths(expr);
-        if match_call_path(
+        if !call_path.is_empty() {
-            &dealias_call_path(call_path, import_aliases),
+            if match_call_path(
-            module,
+                &dealias_call_path(call_path, import_aliases),
-            member,
+                module,
-            from_imports,
+                member,
-        ) {
+                from_imports,
-            return true;
+            ) {
+                return true;
+            }
         }
-    }
+        false
+    })
+}
 
+/// Call `func` over every `Expr` in `expr`, returning `true` if any expression
+/// returns `true`..
+pub fn any_over_expr<F>(expr: &Expr, func: &F) -> bool
+where
+    F: Fn(&Expr) -> bool,
+{
+    if func(expr) {
+        return true;
+    }
     match &expr.node {
-        ExprKind::BoolOp { values, .. } => values
+        ExprKind::BoolOp { values, .. } | ExprKind::JoinedStr { values } => {
-            .iter()
+            values.iter().any(|expr| any_over_expr(expr, func))
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
+        }
         ExprKind::NamedExpr { target, value } => {
-            contains_call_path(target, module, member, import_aliases, from_imports)
+            any_over_expr(target, func) || any_over_expr(value, func)
-                || contains_call_path(value, module, member, import_aliases, from_imports)
         }
         ExprKind::BinOp { left, right, .. } => {
-            contains_call_path(left, module, member, import_aliases, from_imports)
+            any_over_expr(left, func) || any_over_expr(right, func)
-                || contains_call_path(right, module, member, import_aliases, from_imports)
-        }
-        ExprKind::UnaryOp { operand, .. } => {
-            contains_call_path(operand, module, member, import_aliases, from_imports)
-        }
-        ExprKind::Lambda { body, .. } => {
-            contains_call_path(body, module, member, import_aliases, from_imports)
         }
+        ExprKind::UnaryOp { operand, .. } => any_over_expr(operand, func),
+        ExprKind::Lambda { body, .. } => any_over_expr(body, func),
         ExprKind::IfExp { test, body, orelse } => {
-            contains_call_path(test, module, member, import_aliases, from_imports)
+            any_over_expr(test, func) || any_over_expr(body, func) || any_over_expr(orelse, func)
-                || contains_call_path(body, module, member, import_aliases, from_imports)
-                || contains_call_path(orelse, module, member, import_aliases, from_imports)
         }
         ExprKind::Dict { keys, values } => values
             .iter()
             .chain(keys.iter())
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
+            .any(|expr| any_over_expr(expr, func)),
-        ExprKind::Set { elts } => elts
+        ExprKind::Set { elts } | ExprKind::List { elts, .. } | ExprKind::Tuple { elts, .. } => {
-            .iter()
+            elts.iter().any(|expr| any_over_expr(expr, func))
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
-        ExprKind::ListComp { elt, generators } => {
-            contains_call_path(elt, module, member, import_aliases, from_imports)
-                || generators.iter().any(|generator| {
-                    contains_call_path(
-                        &generator.target,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || contains_call_path(
-                        &generator.iter,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || generator.ifs.iter().any(|expr| {
-                        contains_call_path(expr, module, member, import_aliases, from_imports)
-                    })
-                })
         }
-        ExprKind::SetComp { elt, generators } => {
+        ExprKind::ListComp { elt, generators }
-            contains_call_path(elt, module, member, import_aliases, from_imports)
+        | ExprKind::SetComp { elt, generators }
+        | ExprKind::GeneratorExp { elt, generators } => {
+            any_over_expr(elt, func)
                 || generators.iter().any(|generator| {
-                    contains_call_path(
+                    any_over_expr(&generator.target, func)
-                        &generator.target,
+                        || any_over_expr(&generator.iter, func)
-                        module,
+                        || generator.ifs.iter().any(|expr| any_over_expr(expr, func))
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || contains_call_path(
-                        &generator.iter,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || generator.ifs.iter().any(|expr| {
-                        contains_call_path(expr, module, member, import_aliases, from_imports)
-                    })
                 })
         }
         ExprKind::DictComp {
@@ -274,123 +249,58 @@ pub fn contains_call_path(
             value,
             generators,
         } => {
-            contains_call_path(key, module, member, import_aliases, from_imports)
+            any_over_expr(key, func)
-                || contains_call_path(value, module, member, import_aliases, from_imports)
+                || any_over_expr(value, func)
                 || generators.iter().any(|generator| {
-                    contains_call_path(
+                    any_over_expr(&generator.target, func)
-                        &generator.target,
+                        || any_over_expr(&generator.iter, func)
-                        module,
+                        || generator.ifs.iter().any(|expr| any_over_expr(expr, func))
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || contains_call_path(
-                        &generator.iter,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || generator.ifs.iter().any(|expr| {
-                        contains_call_path(expr, module, member, import_aliases, from_imports)
-                    })
                 })
         }
-        ExprKind::GeneratorExp { elt, generators } => {
+        ExprKind::Await { value }
-            contains_call_path(elt, module, member, import_aliases, from_imports) || {
+        | ExprKind::YieldFrom { value }
-                generators.iter().any(|generator| {
+        | ExprKind::Attribute { value, .. }
-                    contains_call_path(
+        | ExprKind::Starred { value, .. } => any_over_expr(value, func),
-                        &generator.target,
+        ExprKind::Yield { value } => value
-                        module,
+            .as_ref()
-                        member,
+            .map_or(false, |value| any_over_expr(value, func)),
-                        import_aliases,
-                        from_imports,
-                    ) || contains_call_path(
-                        &generator.iter,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    ) || generator.ifs.iter().any(|expr| {
-                        contains_call_path(expr, module, member, import_aliases, from_imports)
-                    })
-                })
-            }
-        }
-        ExprKind::Await { value } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
-        }
-        ExprKind::Yield { value } => value.as_ref().map_or(false, |value| {
-            contains_call_path(value, module, member, import_aliases, from_imports)
-        }),
-        ExprKind::YieldFrom { value } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
-        }
         ExprKind::Compare {
-            left,
+            left, comparators, ..
-            ops: _,
+        } => any_over_expr(left, func) || comparators.iter().any(|expr| any_over_expr(expr, func)),
-            comparators,
-        } => {
-            contains_call_path(left, module, member, import_aliases, from_imports)
-                || comparators.iter().any(|expr| {
-                    contains_call_path(expr, module, member, import_aliases, from_imports)
-                })
-        }
         ExprKind::Call {
-            func,
+            func: call_func,
             args,
             keywords,
         } => {
-            contains_call_path(func, module, member, import_aliases, from_imports)
+            any_over_expr(call_func, func)
-                || args.iter().any(|expr| {
+                || args.iter().any(|expr| any_over_expr(expr, func))
-                    contains_call_path(expr, module, member, import_aliases, from_imports)
+                || keywords
-                })
+                    .iter()
-                || keywords.iter().any(|keyword| {
+                    .any(|keyword| any_over_expr(&keyword.node.value, func))
-                    contains_call_path(
-                        &keyword.node.value,
-                        module,
-                        member,
-                        import_aliases,
-                        from_imports,
-                    )
-                })
         }
         ExprKind::FormattedValue {
             value, format_spec, ..
         } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
+            any_over_expr(value, func)
-                || format_spec.as_ref().map_or(false, |value| {
+                || format_spec
-                    contains_call_path(value, module, member, import_aliases, from_imports)
+                    .as_ref()
-                })
+                    .map_or(false, |value| any_over_expr(value, func))
-        }
-        ExprKind::JoinedStr { values } => values
-            .iter()
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
-        ExprKind::Constant { .. } => false,
-        ExprKind::Attribute { value, .. } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
         }
         ExprKind::Subscript { value, slice, .. } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
+            any_over_expr(value, func) || any_over_expr(slice, func)
-                || contains_call_path(slice, module, member, import_aliases, from_imports)
         }
-        ExprKind::Starred { value, ctx: _ } => {
-            contains_call_path(value, module, member, import_aliases, from_imports)
-        }
-        ExprKind::Name { .. } => false,
-        ExprKind::List { elts, .. } => elts
-            .iter()
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
-        ExprKind::Tuple { elts, .. } => elts
-            .iter()
-            .any(|expr| contains_call_path(expr, module, member, import_aliases, from_imports)),
         ExprKind::Slice { lower, upper, step } => {
-            lower.as_ref().map_or(false, |value| {
+            lower
-                contains_call_path(value, module, member, import_aliases, from_imports)
+                .as_ref()
-            }) || upper.as_ref().map_or(false, |value| {
+                .map_or(false, |value| any_over_expr(value, func))
-                contains_call_path(value, module, member, import_aliases, from_imports)
+                || upper
-            }) || step.as_ref().map_or(false, |value| {
+                    .as_ref()
-                contains_call_path(value, module, member, import_aliases, from_imports)
+                    .map_or(false, |value| any_over_expr(value, func))
-            })
+                || step
+                    .as_ref()
+                    .map_or(false, |value| any_over_expr(value, func))
         }
+        ExprKind::Name { .. } | ExprKind::Constant { .. } => false,
     }
 }
 
@@ -406,12 +316,12 @@ pub fn is_assignment_to_a_dunder(stmt: &Stmt) -> bool {
                 return false;
             }
             match &targets[0].node {
-                ExprKind::Name { id, ctx: _ } => DUNDER_REGEX.is_match(id),
+                ExprKind::Name { id, .. } => DUNDER_REGEX.is_match(id),
                 _ => false,
             }
         }
         StmtKind::AnnAssign { target, .. } => match &target.node {
-            ExprKind::Name { id, ctx: _ } => DUNDER_REGEX.is_match(id),
+            ExprKind::Name { id, .. } => DUNDER_REGEX.is_match(id),
             _ => false,
         },
         _ => false,

src/ast/operations.rs

@@ -4,6 +4,7 @@ use rustpython_parser::ast::{Constant, Expr, ExprKind, Stmt, StmtKind};
 use rustpython_parser::lexer;
 use rustpython_parser::lexer::Tok;
 
+use crate::ast::helpers::any_over_expr;
 use crate::ast::types::{Binding, BindingKind, Scope};
 use crate::ast::visitor;
 use crate::ast::visitor::Visitor;
@@ -129,76 +130,13 @@ pub fn in_nested_block<'a>(mut parents: impl Iterator<Item = &'a Stmt>) -> bool
     })
 }
 
-/// Returns `true` if `parent` contains `child`.
-fn contains(parent: &Expr, child: &Expr) -> bool {
-    match &parent.node {
-        ExprKind::BoolOp { values, .. } => values.iter().any(|parent| contains(parent, child)),
-        ExprKind::NamedExpr { target, value } => contains(target, child) || contains(value, child),
-        ExprKind::BinOp { left, right, .. } => contains(left, child) || contains(right, child),
-        ExprKind::UnaryOp { operand, .. } => contains(operand, child),
-        ExprKind::Lambda { body, .. } => contains(body, child),
-        ExprKind::IfExp { test, body, orelse } => {
-            contains(test, child) || contains(body, child) || contains(orelse, child)
-        }
-        ExprKind::Dict { keys, values } => keys
-            .iter()
-            .chain(values.iter())
-            .any(|parent| contains(parent, child)),
-        ExprKind::Set { elts } => elts.iter().any(|parent| contains(parent, child)),
-        ExprKind::ListComp { elt, .. } => contains(elt, child),
-        ExprKind::SetComp { elt, .. } => contains(elt, child),
-        ExprKind::DictComp { key, value, .. } => contains(key, child) || contains(value, child),
-        ExprKind::GeneratorExp { elt, .. } => contains(elt, child),
-        ExprKind::Await { value } => contains(value, child),
-        ExprKind::Yield { value } => value.as_ref().map_or(false, |value| contains(value, child)),
-        ExprKind::YieldFrom { value } => contains(value, child),
-        ExprKind::Compare {
-            left, comparators, ..
-        } => contains(left, child) || comparators.iter().any(|parent| contains(parent, child)),
-        ExprKind::Call {
-            func,
-            args,
-            keywords,
-        } => {
-            contains(func, child)
-                || args.iter().any(|parent| contains(parent, child))
-                || keywords
-                    .iter()
-                    .any(|keyword| contains(&keyword.node.value, child))
-        }
-        ExprKind::FormattedValue {
-            value, format_spec, ..
-        } => {
-            contains(value, child)
-                || format_spec
-                    .as_ref()
-                    .map_or(false, |value| contains(value, child))
-        }
-        ExprKind::JoinedStr { values } => values.iter().any(|parent| contains(parent, child)),
-        ExprKind::Constant { .. } => false,
-        ExprKind::Attribute { value, .. } => contains(value, child),
-        ExprKind::Subscript { value, slice, .. } => {
-            contains(value, child) || contains(slice, child)
-        }
-        ExprKind::Starred { value, .. } => contains(value, child),
-        ExprKind::Name { .. } => parent == child,
-        ExprKind::List { elts, .. } => elts.iter().any(|parent| contains(parent, child)),
-        ExprKind::Tuple { elts, .. } => elts.iter().any(|parent| contains(parent, child)),
-        ExprKind::Slice { lower, upper, step } => {
-            lower.as_ref().map_or(false, |value| contains(value, child))
-                || upper.as_ref().map_or(false, |value| contains(value, child))
-                || step.as_ref().map_or(false, |value| contains(value, child))
-        }
-    }
-}
-
 /// Check if a node represents an unpacking assignment.
 pub fn is_unpacking_assignment(parent: &Stmt, child: &Expr) -> bool {
     match &parent.node {
         StmtKind::With { items, .. } => items.iter().any(|item| {
             if let Some(optional_vars) = &item.optional_vars {
                 if matches!(optional_vars.node, ExprKind::Tuple { .. }) {
-                    if contains(optional_vars, child) {
+                    if any_over_expr(optional_vars, &|expr| expr == child) {
                         return true;
                     }
                 }
@@ -227,7 +165,7 @@ pub fn is_unpacking_assignment(parent: &Stmt, child: &Expr) -> bool {
                     matches!(
                         item.node,
                         ExprKind::Set { .. } | ExprKind::List { .. } | ExprKind::Tuple { .. }
-                    ) && contains(item, child)
+                    ) && any_over_expr(item, &|expr| expr == child)
                 });
 
             // If our child is a tuple, and value is not, it's always an unpacking