Add a specialized StatementVisitor (#4349)

crates/ruff_python_ast/src/helpers.rs

@@ -18,8 +18,7 @@ use smallvec::SmallVec;
 use crate::call_path::CallPath;
 use crate::newlines::UniversalNewlineIterator;
 use crate::source_code::{Generator, Indexer, Locator, Stylist};
-use crate::visitor;
-use crate::visitor::Visitor;
+use crate::statement_visitor::{walk_body, walk_stmt, StatementVisitor};
 
 /// Create an `Expr` with default location from an `ExprKind`.
 pub fn create_expr(node: ExprKind) -> Expr {
@@ -816,13 +815,13 @@ pub fn resolve_imported_module_path<'a>(
     Some(Cow::Owned(qualified_path))
 }
 
-/// A [`Visitor`] that collects all `return` statements in a function or method.
+/// A [`StatementVisitor`] that collects all `return` statements in a function or method.
 #[derive(Default)]
 pub struct ReturnStatementVisitor<'a> {
     pub returns: Vec<Option<&'a Expr>>,
 }
 
-impl<'a, 'b> Visitor<'b> for ReturnStatementVisitor<'a>
+impl<'a, 'b> StatementVisitor<'b> for ReturnStatementVisitor<'a>
 where
     'b: 'a,
 {
@@ -832,18 +831,18 @@ where
                 // Don't recurse.
             }
             StmtKind::Return { value } => self.returns.push(value.as_deref()),
-            _ => visitor::walk_stmt(self, stmt),
+            _ => walk_stmt(self, stmt),
         }
     }
 }
 
-/// A [`Visitor`] that collects all `raise` statements in a function or method.
+/// A [`StatementVisitor`] that collects all `raise` statements in a function or method.
 #[derive(Default)]
 pub struct RaiseStatementVisitor<'a> {
     pub raises: Vec<(TextRange, Option<&'a Expr>, Option<&'a Expr>)>,
 }
 
-impl<'a, 'b> Visitor<'b> for RaiseStatementVisitor<'b>
+impl<'a, 'b> StatementVisitor<'b> for RaiseStatementVisitor<'b>
 where
     'b: 'a,
 {
@@ -859,19 +858,19 @@ where
             | StmtKind::Try { .. }
             | StmtKind::TryStar { .. } => {}
             StmtKind::If { body, orelse, .. } => {
-                visitor::walk_body(self, body);
-                visitor::walk_body(self, orelse);
+                walk_body(self, body);
+                walk_body(self, orelse);
             }
             StmtKind::While { body, .. }
             | StmtKind::With { body, .. }
             | StmtKind::AsyncWith { body, .. }
             | StmtKind::For { body, .. }
             | StmtKind::AsyncFor { body, .. } => {
-                visitor::walk_body(self, body);
+                walk_body(self, body);
             }
             StmtKind::Match { cases, .. } => {
                 for case in cases {
-                    visitor::walk_body(self, &case.body);
+                    walk_body(self, &case.body);
                 }
             }
             _ => {}
@@ -884,7 +883,7 @@ struct GlobalStatementVisitor<'a> {
     globals: FxHashMap<&'a str, &'a Stmt>,
 }
 
-impl<'a> Visitor<'a> for GlobalStatementVisitor<'a> {
+impl<'a> StatementVisitor<'a> for GlobalStatementVisitor<'a> {
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
         match &stmt.node {
             StmtKind::Global { names } => {
@@ -897,7 +896,7 @@ impl<'a> Visitor<'a> for GlobalStatementVisitor<'a> {
             | StmtKind::ClassDef { .. } => {
                 // Don't recurse.
             }
-            _ => visitor::walk_stmt(self, stmt),
+            _ => walk_stmt(self, stmt),
         }
     }
 }

crates/ruff_python_ast/src/lib.rs

@@ -8,6 +8,7 @@ pub mod imports;
 pub mod newlines;
 pub mod relocate;
 pub mod source_code;
+pub mod statement_visitor;
 pub mod str;
 pub mod token_kind;
 pub mod types;

crates/ruff_python_ast/src/statement_visitor.rs

@@ -0,0 +1,111 @@
+//! Specialized AST visitor trait and walk functions that only visit statements.
+
+use rustpython_parser::ast::{Excepthandler, ExcepthandlerKind, MatchCase, Stmt, StmtKind};
+
+/// A trait for AST visitors that only need to visit statements.
+pub trait StatementVisitor<'a> {
+    fn visit_body(&mut self, body: &'a [Stmt]) {
+        walk_body(self, body);
+    }
+    fn visit_stmt(&mut self, stmt: &'a Stmt) {
+        walk_stmt(self, stmt);
+    }
+    fn visit_excepthandler(&mut self, excepthandler: &'a Excepthandler) {
+        walk_excepthandler(self, excepthandler);
+    }
+    fn visit_match_case(&mut self, match_case: &'a MatchCase) {
+        walk_match_case(self, match_case);
+    }
+}
+
+pub fn walk_body<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, body: &'a [Stmt]) {
+    for stmt in body {
+        visitor.visit_stmt(stmt);
+    }
+}
+
+pub fn walk_stmt<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
+    match &stmt.node {
+        StmtKind::FunctionDef { body, .. } => {
+            visitor.visit_body(body);
+        }
+        StmtKind::AsyncFunctionDef { body, .. } => {
+            visitor.visit_body(body);
+        }
+        StmtKind::For { body, orelse, .. } => {
+            visitor.visit_body(body);
+            visitor.visit_body(orelse);
+        }
+        StmtKind::ClassDef { body, .. } => {
+            visitor.visit_body(body);
+        }
+        StmtKind::AsyncFor { body, orelse, .. } => {
+            visitor.visit_body(body);
+            visitor.visit_body(orelse);
+        }
+        StmtKind::While { body, orelse, .. } => {
+            visitor.visit_body(body);
+            visitor.visit_body(orelse);
+        }
+        StmtKind::If { body, orelse, .. } => {
+            visitor.visit_body(body);
+            visitor.visit_body(orelse);
+        }
+        StmtKind::With { body, .. } => {
+            visitor.visit_body(body);
+        }
+        StmtKind::AsyncWith { body, .. } => {
+            visitor.visit_body(body);
+        }
+        StmtKind::Match { cases, .. } => {
+            for match_case in cases {
+                visitor.visit_match_case(match_case);
+            }
+        }
+        StmtKind::Try {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+        } => {
+            visitor.visit_body(body);
+            for excepthandler in handlers {
+                visitor.visit_excepthandler(excepthandler);
+            }
+            visitor.visit_body(orelse);
+            visitor.visit_body(finalbody);
+        }
+        StmtKind::TryStar {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+        } => {
+            visitor.visit_body(body);
+            for excepthandler in handlers {
+                visitor.visit_excepthandler(excepthandler);
+            }
+            visitor.visit_body(orelse);
+            visitor.visit_body(finalbody);
+        }
+        _ => {}
+    }
+}
+
+pub fn walk_excepthandler<'a, V: StatementVisitor<'a> + ?Sized>(
+    visitor: &mut V,
+    excepthandler: &'a Excepthandler,
+) {
+    match &excepthandler.node {
+        ExcepthandlerKind::ExceptHandler { body, .. } => {
+            visitor.visit_body(body);
+        }
+    }
+}
+
+pub fn walk_match_case<'a, V: StatementVisitor<'a> + ?Sized>(
+    visitor: &mut V,
+    match_case: &'a MatchCase,
+) {
+    visitor.visit_body(&match_case.body);
+}

crates/ruff_python_ast/src/visitor.rs

@@ -1,9 +1,15 @@
+//! AST visitor trait and walk functions.
+
 use rustpython_parser::ast::{
     Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
     ExcepthandlerKind, Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern,
     PatternKind, Stmt, StmtKind, Unaryop, Withitem,
 };
 
+/// A trait for AST visitors. Visits all nodes in the AST recursively.
+///
+/// Prefer [`crate::statement_visitor::StatementVisitor`] for visitors that only need to visit
+/// statements.
 pub trait Visitor<'a> {
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
         walk_stmt(self, stmt);