⚡️ Speed up function collect_existing_class_names by 351% in PR #1498 (cf-simplify-context-extraction)

codeflash/languages/python/context/code_context_extractor.py

@@ -563,7 +563,28 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M
 
 
 def collect_existing_class_names(tree: ast.Module) -> set[str]:
-    return {node.name for node in ast.walk(tree) if isinstance(node, ast.ClassDef)}
+    class_names = set()
+    stack = list(tree.body)
+
+    while stack:
+        node = stack.pop()
+        if isinstance(node, ast.ClassDef):
+            class_names.add(node.name)
+            stack.extend(node.body)
+        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            stack.extend(node.body)
+        elif isinstance(node, (ast.If, ast.For, ast.While, ast.With)):
+            stack.extend(node.body)
+            if hasattr(node, "orelse"):
+                stack.extend(node.orelse)
+        elif isinstance(node, ast.Try):
+            stack.extend(node.body)
+            stack.extend(node.orelse)
+            stack.extend(node.finalbody)
+            for handler in node.handlers:
+                stack.extend(handler.body)
+
+    return class_names
 
 
 def enrich_testgen_context(code_context: CodeStringsMarkdown, project_root_path: Path) -> CodeStringsMarkdown: