From e13e57e024adc9c695fbad371a7c35e04b86ce63 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@oddbird.net>
Date: Mon, 8 Apr 2024 13:29:38 -0600
Subject: [PATCH] Localize cleanup for FunctionDef and ClassDef (#10837)

## Summary

Came across this code while digging into the semantic model with
@AlexWaygood, and found it confusing because of how it splits
`push_scope` from the paired `pop_scope` (took me a few minutes to even
figure out if/where we were popping the pushed scope). Since this
"cleanup" is already totally split by node type, there doesn't seem to
be any gain in having it as a separate "step" rather than just
incorporating it into the traversal clauses for those node types.

I left the equivalent cleanup step alone for the expression case,
because in that case it is actually generic across several different
node types, and due to the use of the common `visit_generators` utility
there isn't a clear way to keep the pushes and corresponding pops
localized.

Feel free to just reject this if I've missed a good reason for it to
stay this way!

## Test Plan

Tests and clippy.
---
 crates/ruff_linter/src/checkers/ast/mod.rs | 54 ++++++++++------------
 1 file changed, 25 insertions(+), 29 deletions(-)

diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs
index ea845f0074..86b4f54b67 100644
--- a/crates/ruff_linter/src/checkers/ast/mod.rs
+++ b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -571,6 +571,7 @@ impl<'a> Visitor<'a> for Checker<'a> {
         match stmt {
             Stmt::FunctionDef(
                 function_def @ ast::StmtFunctionDef {
+                    name,
                     body,
                     parameters,
                     decorator_list,
@@ -690,9 +691,21 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 if let Some(globals) = Globals::from_body(body) {
                     self.semantic.set_globals(globals);
                 }
+                let scope_id = self.semantic.scope_id;
+                self.analyze.scopes.push(scope_id);
+                self.semantic.pop_scope(); // Function scope
+                self.semantic.pop_definition();
+                self.semantic.pop_scope(); // Type parameter scope
+                self.add_binding(
+                    name,
+                    stmt.identifier(),
+                    BindingKind::FunctionDefinition(scope_id),
+                    BindingFlags::empty(),
+                );
             }
             Stmt::ClassDef(
                 class_def @ ast::StmtClassDef {
+                    name,
                     body,
                     arguments,
                     decorator_list,
@@ -733,6 +746,18 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 // Set the docstring state before visiting the class body.
                 self.docstring_state = DocstringState::Expected;
                 self.visit_body(body);
+
+                let scope_id = self.semantic.scope_id;
+                self.analyze.scopes.push(scope_id);
+                self.semantic.pop_scope(); // Class scope
+                self.semantic.pop_definition();
+                self.semantic.pop_scope(); // Type parameter scope
+                self.add_binding(
+                    name,
+                    stmt.identifier(),
+                    BindingKind::ClassDefinition(scope_id),
+                    BindingFlags::empty(),
+                );
             }
             Stmt::TypeAlias(ast::StmtTypeAlias {
                 range: _,
@@ -889,35 +914,6 @@ impl<'a> Visitor<'a> for Checker<'a> {
         };
 
         // Step 3: Clean-up
-        match stmt {
-            Stmt::FunctionDef(ast::StmtFunctionDef { name, .. }) => {
-                let scope_id = self.semantic.scope_id;
-                self.analyze.scopes.push(scope_id);
-                self.semantic.pop_scope(); // Function scope
-                self.semantic.pop_definition();
-                self.semantic.pop_scope(); // Type parameter scope
-                self.add_binding(
-                    name,
-                    stmt.identifier(),
-                    BindingKind::FunctionDefinition(scope_id),
-                    BindingFlags::empty(),
-                );
-            }
-            Stmt::ClassDef(ast::StmtClassDef { name, .. }) => {
-                let scope_id = self.semantic.scope_id;
-                self.analyze.scopes.push(scope_id);
-                self.semantic.pop_scope(); // Class scope
-                self.semantic.pop_definition();
-                self.semantic.pop_scope(); // Type parameter scope
-                self.add_binding(
-                    name,
-                    stmt.identifier(),
-                    BindingKind::ClassDefinition(scope_id),
-                    BindingFlags::empty(),
-                );
-            }
-            _ => {}
-        }
 
         // Step 4: Analysis
         analyze::statement(stmt, self);