rustpython_ruff_python_ast/
statement_visitor.rs

1//! Specialized AST visitor trait and walk functions that only visit statements.
2
3use crate::{self as ast, ElifElseClause, ExceptHandler, MatchCase, Stmt};
4
5/// A trait for AST visitors that only need to visit statements.
6pub trait StatementVisitor<'a> {
7    fn visit_body(&mut self, body: &'a [Stmt]) {
8        walk_body(self, body);
9    }
10    fn visit_stmt(&mut self, stmt: &'a Stmt) {
11        walk_stmt(self, stmt);
12    }
13    fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
14        walk_except_handler(self, except_handler);
15    }
16    fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
17        walk_elif_else_clause(self, elif_else_clause);
18    }
19    fn visit_match_case(&mut self, match_case: &'a MatchCase) {
20        walk_match_case(self, match_case);
21    }
22}
23
24pub fn walk_body<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, body: &'a [Stmt]) {
25    for stmt in body {
26        visitor.visit_stmt(stmt);
27    }
28}
29
30pub fn walk_stmt<'a, V: StatementVisitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
31    match stmt {
32        Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => {
33            visitor.visit_body(body);
34        }
35        Stmt::For(ast::StmtFor { body, orelse, .. }) => {
36            visitor.visit_body(body);
37            visitor.visit_body(orelse);
38        }
39        Stmt::ClassDef(ast::StmtClassDef { body, .. }) => {
40            visitor.visit_body(body);
41        }
42        Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
43            visitor.visit_body(body);
44            visitor.visit_body(orelse);
45        }
46        Stmt::If(ast::StmtIf {
47            body,
48            elif_else_clauses,
49            ..
50        }) => {
51            visitor.visit_body(body);
52            for clause in elif_else_clauses {
53                visitor.visit_elif_else_clause(clause);
54            }
55        }
56        Stmt::With(ast::StmtWith { body, .. }) => {
57            visitor.visit_body(body);
58        }
59        Stmt::Match(ast::StmtMatch { cases, .. }) => {
60            for match_case in cases {
61                visitor.visit_match_case(match_case);
62            }
63        }
64        Stmt::Try(ast::StmtTry {
65            body,
66            handlers,
67            orelse,
68            finalbody,
69            ..
70        }) => {
71            visitor.visit_body(body);
72            for except_handler in handlers {
73                visitor.visit_except_handler(except_handler);
74            }
75            visitor.visit_body(orelse);
76            visitor.visit_body(finalbody);
77        }
78        _ => {}
79    }
80}
81
82pub fn walk_except_handler<'a, V: StatementVisitor<'a> + ?Sized>(
83    visitor: &mut V,
84    except_handler: &'a ExceptHandler,
85) {
86    match except_handler {
87        ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { body, .. }) => {
88            visitor.visit_body(body);
89        }
90    }
91}
92
93pub fn walk_elif_else_clause<'a, V: StatementVisitor<'a> + ?Sized>(
94    visitor: &mut V,
95    elif_else_clause: &'a ElifElseClause,
96) {
97    visitor.visit_body(&elif_else_clause.body);
98}
99
100pub fn walk_match_case<'a, V: StatementVisitor<'a> + ?Sized>(
101    visitor: &mut V,
102    match_case: &'a MatchCase,
103) {
104    visitor.visit_body(&match_case.body);
105}
rustpython_ruff_python_ast/statement_visitor.rs

rustpython_ruff_python_ast/
statement_visitor.rs
