⚡️ Speed up function _parse_and_collect_imports by 12% in PR #1660 (unstructured-inference)

codeflash/languages/python/context/code_context_extractor.py

@@ -6,7 +6,6 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
 from itertools import chain
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 import libcst as cst
@@ -40,6 +39,8 @@
 )
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from jedi.api.classes import Name
 
     from codeflash.languages.base import DependencyResolver
@@ -566,13 +567,9 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M
         tree = ast.parse(all_code)
     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_names[alias.asname if alias.asname else alias.name] = node.module
-    return tree, imported_names
+    collector = ImportCollector()
+    collector.visit(tree)
+    return tree, collector.imported_names
 
 
 def collect_existing_class_names(tree: ast.Module) -> set[str]:
@@ -922,6 +919,7 @@ def _has_descriptor_like_class_fields(class_node: ast.ClassDef) -> bool:
 
 def _should_use_raw_project_class_context(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> bool:
     start_line = _get_class_start_line(class_node)
+    assert class_node.end_lineno is not None
     class_line_count = class_node.end_lineno - start_line + 1
     is_small = (
         class_line_count <= MAX_RAW_PROJECT_CLASS_LINES and len(class_node.body) <= MAX_RAW_PROJECT_CLASS_BODY_ITEMS
@@ -1490,6 +1488,17 @@ def _maybe_strip_docstring(node: cst.FunctionDef | cst.ClassDef, cfg: PruneConfi
     return node
 
 
+class ImportCollector(ast.NodeVisitor):
+    def __init__(self) -> None:
+        self.imported_names: dict[str, str] = {}
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        if node.module:
+            for alias in node.names:
+                if alias.name != "*":
+                    self.imported_names[alias.asname if alias.asname else alias.name] = node.module
+
+
 @dataclass(frozen=True)
 class PruneConfig:
     defs_with_usages: dict[str, UsageInfo] | None = None