From 52bab329844fe1da121587b7b3c5ca1eb0a47d94 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 21 Sep 2024 14:54:11 -0400
Subject: [PATCH] Skip traversal for non-compound statements

---
 crates/ruff_graph/src/collector.rs | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/crates/ruff_graph/src/collector.rs b/crates/ruff_graph/src/collector.rs
index c9b3a6f9b34cfc..948bc4cd0bc96a 100644
--- a/crates/ruff_graph/src/collector.rs
+++ b/crates/ruff_graph/src/collector.rs
@@ -1,8 +1,8 @@
 use red_knot_python_semantic::ModuleName;
 use ruff_python_ast::visitor::source_order::{
-    walk_expr, walk_module, walk_stmt, SourceOrderVisitor,
+    walk_expr, walk_module, walk_stmt, SourceOrderVisitor, TraversalSignal,
 };
-use ruff_python_ast::{self as ast, Expr, Mod, Stmt};
+use ruff_python_ast::{self as ast, AnyNodeRef, Expr, Mod, Stmt};
 
 /// Collect all imports for a given Python file.
 #[derive(Default, Debug)]
@@ -32,6 +32,26 @@ impl<'a> Collector<'a> {
 }
 
 impl<'ast> SourceOrderVisitor<'ast> for Collector<'_> {
+    fn enter_node(&mut self, node: AnyNodeRef<'ast>) -> TraversalSignal {
+        if self.string_imports {
+            TraversalSignal::Traverse
+        } else if matches!(
+            node,
+            AnyNodeRef::ModModule(_)
+                | AnyNodeRef::StmtFunctionDef(_)
+                | AnyNodeRef::StmtClassDef(_)
+                | AnyNodeRef::StmtWhile(_)
+                | AnyNodeRef::StmtFor(_)
+                | AnyNodeRef::StmtWith(_)
+                | AnyNodeRef::StmtIf(_)
+                | AnyNodeRef::StmtTry(_)
+        ) {
+            TraversalSignal::Traverse
+        } else {
+            TraversalSignal::Skip
+        }
+    }
+
     fn visit_stmt(&mut self, stmt: &'ast Stmt) {
         match stmt {
             Stmt::ImportFrom(ast::StmtImportFrom {