-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Use referencial equality in
traversal helper methods (#13895)
- Loading branch information
1 parent
de4181d
commit e402e27
Showing
4 changed files
with
66 additions
and
65 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,81 +1,81 @@ | ||
| //! Utilities for manually traversing a Python AST. | ||
| use crate::{self as ast, ExceptHandler, Stmt, Suite}; | ||
| use crate::{self as ast, AnyNodeRef, ExceptHandler, Stmt}; | ||
|
|
||
| /// 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> { | ||
| /// Given a [`Stmt`] and its parent, return the [`ast::Suite`] that contains the [`Stmt`]. | ||
| pub fn suite<'a>(stmt: &'a Stmt, parent: &'a Stmt) -> Option<EnclosingSuite<'a>> { | ||
| // TODO: refactor this to work without a parent, ie when `stmt` is at the top level | ||
| match parent { | ||
| Stmt::FunctionDef(ast::StmtFunctionDef { 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::While(ast::StmtWhile { body, orelse, .. }) => { | ||
| if body.contains(stmt) { | ||
| Some(body) | ||
| } else if orelse.contains(stmt) { | ||
| Some(orelse) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => EnclosingSuite::new(body, stmt), | ||
| Stmt::ClassDef(ast::StmtClassDef { body, .. }) => EnclosingSuite::new(body, stmt), | ||
| Stmt::For(ast::StmtFor { body, orelse, .. }) => [body, orelse] | ||
| .iter() | ||
| .find_map(|suite| EnclosingSuite::new(suite, stmt)), | ||
| Stmt::While(ast::StmtWhile { body, orelse, .. }) => [body, orelse] | ||
| .iter() | ||
| .find_map(|suite| EnclosingSuite::new(suite, stmt)), | ||
| 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), | ||
| }) => [body] | ||
| .into_iter() | ||
| .chain(elif_else_clauses.iter().map(|clause| &clause.body)) | ||
| .find_map(|suite| EnclosingSuite::new(suite, stmt)), | ||
| Stmt::With(ast::StmtWith { body, .. }) => EnclosingSuite::new(body, stmt), | ||
| Stmt::Match(ast::StmtMatch { cases, .. }) => cases | ||
| .iter() | ||
| .map(|case| &case.body) | ||
| .find(|body| body.contains(stmt)), | ||
| .find_map(|body| EnclosingSuite::new(body, 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 { | ||
| }) => [body, orelse, finalbody] | ||
| .into_iter() | ||
| .chain( | ||
| handlers | ||
| .iter() | ||
| .filter_map(ExceptHandler::as_except_handler) | ||
| .map(|handler| &handler.body) | ||
| .find(|body| body.contains(stmt)) | ||
| } | ||
| } | ||
| .map(|handler| &handler.body), | ||
| ) | ||
| .find_map(|suite| EnclosingSuite::new(suite, 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(); | ||
| } | ||
| pub struct EnclosingSuite<'a> { | ||
| suite: &'a [Stmt], | ||
| position: usize, | ||
| } | ||
|
|
||
| impl<'a> EnclosingSuite<'a> { | ||
| pub fn new(suite: &'a [Stmt], stmt: &'a Stmt) -> Option<Self> { | ||
| let position = suite | ||
| .iter() | ||
| .position(|sibling| AnyNodeRef::ptr_eq(sibling.into(), stmt.into()))?; | ||
|
|
||
| Some(EnclosingSuite { suite, position }) | ||
| } | ||
|
|
||
| pub fn next_sibling(&self) -> Option<&'a Stmt> { | ||
| self.suite.get(self.position + 1) | ||
| } | ||
|
|
||
| pub fn next_siblings(&self) -> &'a [Stmt] { | ||
| self.suite.get(self.position + 1..).unwrap_or_default() | ||
| } | ||
|
|
||
| pub fn previous_sibling(&self) -> Option<&'a Stmt> { | ||
| self.suite.get(self.position.checked_sub(1)?) | ||
| } | ||
| } | ||
|
|
||
| impl std::ops::Deref for EnclosingSuite<'_> { | ||
| type Target = [Stmt]; | ||
|
|
||
| fn deref(&self) -> &Self::Target { | ||
| self.suite | ||
| } | ||
| None | ||
| } |