Merge pull request #1499 from codeflash-ai/codeflash/optimize-pr1498-2026-02-16T20.49.33

⚡️ Speed up function `_parse_and_collect_imports` by 69% in PR #1498 (`cf-simplify-context-extraction`)

Author: KRRT7

codeflash/languages/python/context/code_context_extractor.py

@@ -553,12 +553,48 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M
     except SyntaxError:
         return None
     imported_names: dict[str, str] = {}
-    for node in ast.walk(tree):
-        if isinstance(node, ast.ImportFrom) and node.module:
-            for alias in node.names:
-                if alias.name != "*":
-                    imported_name = alias.asname if alias.asname else alias.name
-                    imported_names[imported_name] = node.module
+
+    # Directly iterate over the module body and nested structures instead of ast.walk
+    # This avoids traversing every single node in the tree
+    def collect_imports(nodes: list[ast.stmt]) -> None:
+        for node in nodes:
+            if isinstance(node, ast.ImportFrom) and node.module:
+                for alias in node.names:
+                    if alias.name != "*":
+                        imported_name = alias.asname if alias.asname else alias.name
+                        imported_names[imported_name] = node.module
+            # Recursively check nested structures (function defs, class defs, if statements, etc.)
+            elif isinstance(
+                node,
+                (
+                    ast.FunctionDef,
+                    ast.AsyncFunctionDef,
+                    ast.ClassDef,
+                    ast.If,
+                    ast.For,
+                    ast.AsyncFor,
+                    ast.While,
+                    ast.With,
+                    ast.AsyncWith,
+                    ast.Try,
+                    ast.ExceptHandler,
+                ),
+            ):
+                if hasattr(node, "body"):
+                    collect_imports(node.body)
+                if hasattr(node, "orelse"):
+                    collect_imports(node.orelse)
+                if hasattr(node, "finalbody"):
+                    collect_imports(node.finalbody)
+                if hasattr(node, "handlers"):
+                    for handler in node.handlers:
+                        collect_imports(handler.body)
+            # Handle match/case statements (Python 3.10+)
+            elif hasattr(ast, "Match") and isinstance(node, ast.Match):
+                for case in node.cases:
+                    collect_imports(case.body)
+
+    collect_imports(tree.body)
     return tree, imported_names