mirror of
https://github.com/astral-sh/ruff.git
synced 2025-08-22 11:25:26 +00:00
2d505e2b04
## Summary
The `SemanticModel` currently stores the "body" of a given `Suite`,
along with the current statement index. This is used to support "next
sibling" queries, but we only use this in exactly one place -- the rule
that simplifies constructs like this to `any` or `all`:
```python
for x in y:
if x == 0:
return True
return False
```
Instead of tracking the state, we can just do a (slightly more
expensive) traversal, by finding the node within its parent and
returning the next node in the body.
Note that we'll only have to do this extremely rarely -- namely, for
functions that contain something like:
```python
for x in y:
if x == 0:
return True
```
113 lines
3.6 KiB
Rust
113 lines
3.6 KiB
Rust
//! Utilities for manually traversing a Python AST.
|
|
use rustpython_ast::{ExceptHandler, Stmt, Suite};
|
|
use rustpython_parser::ast;
|
|
|
|
/// Given a [`Stmt`] and its parent, return the [`Suite`] that contains the [`Stmt`].
|
|
pub fn suite<'a>(stmt: &'a Stmt, parent: &'a Stmt) -> Option<&'a Suite> {
|
|
match parent {
|
|
Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => Some(body),
|
|
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. }) => Some(body),
|
|
Stmt::ClassDef(ast::StmtClassDef { body, .. }) => Some(body),
|
|
Stmt::For(ast::StmtFor { body, orelse, .. }) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else if orelse.contains(stmt) {
|
|
Some(orelse)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Stmt::AsyncFor(ast::StmtAsyncFor { body, orelse, .. }) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else if orelse.contains(stmt) {
|
|
Some(orelse)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else if orelse.contains(stmt) {
|
|
Some(orelse)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Stmt::If(ast::StmtIf {
|
|
body,
|
|
elif_else_clauses,
|
|
..
|
|
}) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else {
|
|
elif_else_clauses
|
|
.iter()
|
|
.map(|elif_else_clause| &elif_else_clause.body)
|
|
.find(|body| body.contains(stmt))
|
|
}
|
|
}
|
|
Stmt::With(ast::StmtWith { body, .. }) => Some(body),
|
|
Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => Some(body),
|
|
Stmt::Match(ast::StmtMatch { cases, .. }) => cases
|
|
.iter()
|
|
.map(|case| &case.body)
|
|
.find(|body| body.contains(stmt)),
|
|
Stmt::Try(ast::StmtTry {
|
|
body,
|
|
handlers,
|
|
orelse,
|
|
finalbody,
|
|
..
|
|
}) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else if orelse.contains(stmt) {
|
|
Some(orelse)
|
|
} else if finalbody.contains(stmt) {
|
|
Some(finalbody)
|
|
} else {
|
|
handlers
|
|
.iter()
|
|
.filter_map(ExceptHandler::as_except_handler)
|
|
.map(|handler| &handler.body)
|
|
.find(|body| body.contains(stmt))
|
|
}
|
|
}
|
|
Stmt::TryStar(ast::StmtTryStar {
|
|
body,
|
|
handlers,
|
|
orelse,
|
|
finalbody,
|
|
..
|
|
}) => {
|
|
if body.contains(stmt) {
|
|
Some(body)
|
|
} else if orelse.contains(stmt) {
|
|
Some(orelse)
|
|
} else if finalbody.contains(stmt) {
|
|
Some(finalbody)
|
|
} else {
|
|
handlers
|
|
.iter()
|
|
.filter_map(ExceptHandler::as_except_handler)
|
|
.map(|handler| &handler.body)
|
|
.find(|body| body.contains(stmt))
|
|
}
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// Given a [`Stmt`] and its containing [`Suite`], return the next [`Stmt`] in the [`Suite`].
|
|
pub fn next_sibling<'a>(stmt: &'a Stmt, suite: &'a Suite) -> Option<&'a Stmt> {
|
|
let mut iter = suite.iter();
|
|
while let Some(sibling) = iter.next() {
|
|
if sibling == stmt {
|
|
return iter.next();
|
|
}
|
|
}
|
|
None
|
|
}
|