From 2142bf6141db67adb3845837f89af9373f3f8482 Mon Sep 17 00:00:00 2001
From: James Berry <jgbcode@gmail.com>
Date: Thu, 22 Jun 2023 23:55:42 -0400
Subject: [PATCH] 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.
---
 crates/ruff_python_ast/src/visitor.rs          | 12 ++++++++++--
 crates/ruff_python_ast/src/visitor/preorder.rs | 15 +++++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/crates/ruff_python_ast/src/visitor.rs b/crates/ruff_python_ast/src/visitor.rs
index 766ed29441..9cf9d7786b 100644
--- a/crates/ruff_python_ast/src/visitor.rs
+++ b/crates/ruff_python_ast/src/visitor.rs
@@ -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 {
diff --git a/crates/ruff_python_ast/src/visitor/preorder.rs b/crates/ruff_python_ast/src/visitor/preorder.rs
index 544a0c8402..8bf9deabfe 100644
--- a/crates/ruff_python_ast/src/visitor/preorder.rs
+++ b/crates/ruff_python_ast/src/visitor/preorder.rs
@@ -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,