Fix annotation and format spec visitors (#5324)

## Summary

The `Visitor` and `preorder::Visitor` traits provide some convenience
functions, `visit_annotation` and `visit_format_spec`, for handling
annotation and format spec expressions respectively. Both of these
functions accept an `&Expr` and have a default implementation which
delegates to `walk_expr`. The problem with this approach is that any
custom handling done in `visit_expr` will be skipped for annotations and
format specs. Instead, to capture any custom logic implemented in
`visit_expr`, both of these function's default implementations should
delegate to `visit_expr` instead of `walk_expr`.

## Example

Consider the below `Visitor` implementation:
```rust
impl<'a> Visitor<'a> for Example<'a> {
    fn visit_expr(&mut self, expr: &'a Expr) {
        match expr {
            Expr::Name(ExprName { id, .. }) => println!("Visiting {:?}", id),
            _ => walk_expr(self, expr),
        }
    }
}
```

Run on the following Python snippet:
```python
a: b
```

I would expect such a visitor to print the following:
```
Visiting b
Visiting a
```

But it instead prints the following:
```
Visiting a
```

Our custom `visit_expr` handler is not invoked for the annotation.

## Test Plan

Tests added in #5271 caught this behavior.
This commit is contained in:
James Berry 2023-06-22 23:55:42 -04:00 committed by GitHub
parent 1cf307c34c
commit 2142bf6141
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 4 deletions

View file

@ -20,7 +20,7 @@ pub trait Visitor<'a> {
walk_stmt(self, stmt);
}
fn visit_annotation(&mut self, expr: &'a Expr) {
walk_expr(self, expr);
walk_annotation(self, expr);
}
fn visit_decorator(&mut self, decorator: &'a Decorator) {
walk_decorator(self, decorator);
@ -53,7 +53,7 @@ pub trait Visitor<'a> {
walk_except_handler(self, except_handler);
}
fn visit_format_spec(&mut self, format_spec: &'a Expr) {
walk_expr(self, format_spec);
walk_format_spec(self, format_spec);
}
fn visit_arguments(&mut self, arguments: &'a Arguments) {
walk_arguments(self, arguments);
@ -322,6 +322,10 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
}
}
pub fn walk_annotation<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(expr);
}
pub fn walk_decorator<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, decorator: &'a Decorator) {
visitor.visit_expr(&decorator.expression);
}
@ -602,6 +606,10 @@ pub fn walk_except_handler<'a, V: Visitor<'a> + ?Sized>(
}
}
pub fn walk_format_spec<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, format_spec: &'a Expr) {
visitor.visit_expr(format_spec);
}
pub fn walk_arguments<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, arguments: &'a Arguments) {
// Defaults are evaluated before annotations.
for arg in &arguments.posonlyargs {

View file

@ -11,7 +11,7 @@ pub trait PreorderVisitor<'a> {
}
fn visit_annotation(&mut self, expr: &'a Expr) {
walk_expr(self, expr);
walk_annotation(self, expr);
}
fn visit_expr(&mut self, expr: &'a Expr) {
@ -51,7 +51,7 @@ pub trait PreorderVisitor<'a> {
}
fn visit_format_spec(&mut self, format_spec: &'a Expr) {
walk_expr(self, format_spec);
walk_format_spec(self, format_spec);
}
fn visit_arguments(&mut self, arguments: &'a Arguments) {
@ -395,6 +395,10 @@ where
}
}
pub fn walk_annotation<'a, V: PreorderVisitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(expr);
}
pub fn walk_decorator<'a, V>(visitor: &mut V, decorator: &'a Decorator)
where
V: PreorderVisitor<'a> + ?Sized,
@ -726,6 +730,13 @@ where
}
}
pub fn walk_format_spec<'a, V: PreorderVisitor<'a> + ?Sized>(
visitor: &mut V,
format_spec: &'a Expr,
) {
visitor.visit_expr(format_spec);
}
pub fn walk_arguments<'a, V>(visitor: &mut V, arguments: &'a Arguments)
where
V: PreorderVisitor<'a> + ?Sized,