feat: add reference graph for Python

codeflash/cli_cmds/console.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import logging
+from collections import deque
 from contextlib import contextmanager
 from itertools import cycle
 from typing import TYPE_CHECKING, Optional
 
 from rich.console import Console
 from rich.logging import RichHandler
+from rich.panel import Panel
 from rich.progress import (
     BarColumn,
     MofNCompleteColumn,
@@ -24,10 +26,13 @@
 from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Callable, Generator
+    from pathlib import Path
 
     from rich.progress import TaskID
 
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+    from codeflash.languages.base import DependencyResolver, IndexResult
     from codeflash.lsp.lsp_message import LspMessage
 
 DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG
@@ -196,3 +201,151 @@ def test_files_progress_bar(total: int, description: str) -> Generator[tuple[Pro
     ) as progress:
         task_id = progress.add_task(description, total=total)
         yield progress, task_id
+
+
+MAX_TREE_ENTRIES = 8
+
+
+@contextmanager
+def call_graph_live_display(
+    total: int, project_root: Path | None = None
+) -> Generator[Callable[[IndexResult], None], None, None]:
+    from rich.console import Group
+    from rich.live import Live
+    from rich.panel import Panel
+    from rich.text import Text
+    from rich.tree import Tree
+
+    if is_LSP_enabled():
+        lsp_log(LspTextMessage(text="Building call graph", takes_time=True))
+        yield lambda _: None
+        return
+
+    progress = Progress(
+        SpinnerColumn(next(spinners)),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(complete_style="cyan", finished_style="green", pulse_style="yellow"),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        auto_refresh=False,
+    )
+    task_id = progress.add_task("Analyzing files", total=total)
+
+    results: deque[IndexResult] = deque(maxlen=MAX_TREE_ENTRIES)
+    stats = {"indexed": 0, "cached": 0, "edges": 0, "external": 0, "errors": 0}
+
+    tree = Tree("[bold]Recent Files[/bold]")
+    stats_text = Text("0 calls found", style="dim")
+    panel = Panel(
+        Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan"
+    )
+
+    def create_tree_node(result: IndexResult) -> Tree:
+        if project_root:
+            try:
+                name = str(result.file_path.resolve().relative_to(project_root.resolve()))
+            except ValueError:
+                name = f"{result.file_path.parent.name}/{result.file_path.name}"
+        else:
+            name = f"{result.file_path.parent.name}/{result.file_path.name}"
+
+        if result.error:
+            return Tree(f"[red]{name}  (error)[/red]")
+
+        if result.cached:
+            return Tree(f"[dim]{name}  (cached)[/dim]")
+
+        local_edges = result.num_edges - result.cross_file_edges
+        edge_info = []
+
+        if local_edges:
+            edge_info.append(f"{local_edges} calls in same file")
+        if result.cross_file_edges:
+            edge_info.append(f"{result.cross_file_edges} calls from other modules")
+
+        label = ", ".join(edge_info) if edge_info else "no calls"
+        return Tree(f"[cyan]{name}[/cyan]  [dim]{label}[/dim]")
+
+    def refresh_display() -> None:
+        tree.children = [create_tree_node(r) for r in results]
+        tree.children.extend([Tree(" ")] * (MAX_TREE_ENTRIES - len(results)))
+
+        # Update stats
+        stat_parts = []
+        if stats["indexed"]:
+            stat_parts.append(f"{stats['indexed']} files analyzed")
+        if stats["cached"]:
+            stat_parts.append(f"{stats['cached']} cached")
+        if stats["errors"]:
+            stat_parts.append(f"{stats['errors']} errors")
+        stat_parts.append(f"{stats['edges']} calls found")
+        if stats["external"]:
+            stat_parts.append(f"{stats['external']} cross-file calls")
+
+        stats_text.truncate(0)
+        stats_text.append(" · ".join(stat_parts), style="dim")
+
+    batch: list[IndexResult] = []
+
+    def process_batch() -> None:
+        for result in batch:
+            results.append(result)
+
+            if result.error:
+                stats["errors"] += 1
+            elif result.cached:
+                stats["cached"] += 1
+            else:
+                stats["indexed"] += 1
+                stats["edges"] += result.num_edges
+                stats["external"] += result.cross_file_edges
+
+            progress.advance(task_id)
+
+        batch.clear()
+        refresh_display()
+        live.refresh()
+
+    def update(result: IndexResult) -> None:
+        batch.append(result)
+        if len(batch) >= 8:
+            process_batch()
+
+    with Live(panel, console=console, transient=False, auto_refresh=False) as live:
+        yield update
+        if batch:
+            process_batch()
+
+
+def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None:
+    total_functions = sum(map(len, file_to_funcs.values()))
+    if not total_functions:
+        return
+
+    # Build the mapping expected by the dependency resolver
+    file_items = file_to_funcs.items()
+    mapping = {file_path: {func.qualified_name for func in funcs} for file_path, funcs in file_items}
+
+    callee_counts = call_graph.count_callees_per_function(mapping)
+
+    # Use built-in sum for C-level loops to reduce Python overhead
+    total_callees = sum(callee_counts.values())
+    with_context = sum(1 for count in callee_counts.values() if count > 0)
+
+    leaf_functions = total_functions - with_context
+    avg_callees = total_callees / total_functions
+
+    function_label = "function" if total_functions == 1 else "functions"
+
+    summary = (
+        f"{total_functions} {function_label} ready for optimization\n"
+        f"Uses other functions: {with_context} · "
+        f"Standalone: {leaf_functions}"
+    )
+
+    if is_LSP_enabled():
+        lsp_log(LspTextMessage(text=summary))
+        return
+
+    console.print(Panel(summary, title="Call Graph Summary", border_style="cyan"))

codeflash/code_utils/code_replacer.py

@@ -871,8 +871,7 @@ def replace_optimized_code(
                     [
                         callee.qualified_name
                         for callee in code_context.helper_functions
-                        if callee.file_path == module_path
-                        and (callee.jedi_definition is None or callee.jedi_definition.type != "class")
+                        if callee.file_path == module_path and callee.definition_type != "class"
                     ]
                 ),
                 candidate.source_code,

codeflash/code_utils/compat.py

@@ -2,46 +2,16 @@
 import sys
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from platformdirs import user_config_dir
 
-if TYPE_CHECKING:
-    codeflash_temp_dir: Path
-    codeflash_cache_dir: Path
-    codeflash_cache_db: Path
+LF: str = os.linesep
+IS_POSIX: bool = os.name != "nt"
+SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix()
 
+codeflash_cache_dir: Path = Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True))
 
-class Compat:
-    # os-independent newline
-    LF: str = os.linesep
+codeflash_temp_dir: Path = Path(tempfile.gettempdir()) / "codeflash"
+codeflash_temp_dir.mkdir(parents=True, exist_ok=True)
 
-    SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix()
-
-    IS_POSIX: bool = os.name != "nt"
-
-    @property
-    def codeflash_cache_dir(self) -> Path:
-        return Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True))
-
-    @property
-    def codeflash_temp_dir(self) -> Path:
-        temp_dir = Path(tempfile.gettempdir()) / "codeflash"
-        if not temp_dir.exists():
-            temp_dir.mkdir(parents=True, exist_ok=True)
-        return temp_dir
-
-    @property
-    def codeflash_cache_db(self) -> Path:
-        return self.codeflash_cache_dir / "codeflash_cache.db"
-
-
-_compat = Compat()
-
-
-codeflash_temp_dir = _compat.codeflash_temp_dir
-codeflash_cache_dir = _compat.codeflash_cache_dir
-codeflash_cache_db = _compat.codeflash_cache_db
-LF = _compat.LF
-SAFE_SYS_EXECUTABLE = _compat.SAFE_SYS_EXECUTABLE
-IS_POSIX = _compat.IS_POSIX
+codeflash_cache_db: Path = codeflash_cache_dir / "codeflash_cache.db"

codeflash/code_utils/config_consts.py

@@ -4,8 +4,8 @@
 from typing import Any, Union
 
 MAX_TEST_RUN_ITERATIONS = 5
-OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 48000
-TESTGEN_CONTEXT_TOKEN_LIMIT = 48000
+OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 64000
+TESTGEN_CONTEXT_TOKEN_LIMIT = 64000
 INDIVIDUAL_TESTCASE_TIMEOUT = 15
 MAX_FUNCTION_TEST_SECONDS = 60
 MIN_IMPROVEMENT_THRESHOLD = 0.05

codeflash/discovery/functions_to_optimize.py

@@ -114,32 +114,30 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
             )
 
 
-class FunctionWithReturnStatement(ast.NodeVisitor):
-    def __init__(self, file_path: Path) -> None:
-        self.functions: list[FunctionToOptimize] = []
-        self.ast_path: list[FunctionParent] = []
-        self.file_path: Path = file_path
-
-    def visit_FunctionDef(self, node: FunctionDef) -> None:
-        if function_has_return_statement(node) and not function_is_a_property(node):
-            self.functions.append(
-                FunctionToOptimize(function_name=node.name, file_path=self.file_path, parents=self.ast_path[:])
-            )
-
-    def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None:
-        if function_has_return_statement(node) and not function_is_a_property(node):
-            self.functions.append(
-                FunctionToOptimize(
-                    function_name=node.name, file_path=self.file_path, parents=self.ast_path[:], is_async=True
+def find_functions_with_return_statement(ast_module: ast.Module, file_path: Path) -> list[FunctionToOptimize]:
+    results: list[FunctionToOptimize] = []
+    # (node, parent_path) — iterative DFS avoids RecursionError on deeply nested ASTs
+    stack: list[tuple[ast.AST, list[FunctionParent]]] = [(ast_module, [])]
+    while stack:
+        node, ast_path = stack.pop()
+        if isinstance(node, (FunctionDef, AsyncFunctionDef)):
+            if function_has_return_statement(node) and not function_is_a_property(node):
+                results.append(
+                    FunctionToOptimize(
+                        function_name=node.name,
+                        file_path=file_path,
+                        parents=ast_path[:],
+                        is_async=isinstance(node, AsyncFunctionDef),
+                    )
                 )
-            )
-
-    def generic_visit(self, node: ast.AST) -> None:
-        if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)):
-            self.ast_path.append(FunctionParent(node.name, node.__class__.__name__))
-        super().generic_visit(node)
-        if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)):
-            self.ast_path.pop()
+            # Don't recurse into function bodies (matches original visitor behaviour)
+            continue
+        child_path = (
+            [*ast_path, FunctionParent(node.name, node.__class__.__name__)] if isinstance(node, ClassDef) else ast_path
+        )
+        for child in reversed(list(ast.iter_child_nodes(node))):
+            stack.append((child, child_path))
+    return results
 
 
 # =============================================================================
@@ -265,9 +263,7 @@ def _find_all_functions_in_python_file(file_path: Path) -> dict[Path, list[Funct
             if DEBUG_MODE:
                 logger.exception(e)
             return functions
-        function_name_visitor = FunctionWithReturnStatement(file_path)
-        function_name_visitor.visit(ast_module)
-        functions[file_path] = function_name_visitor.functions
+        functions[file_path] = find_functions_with_return_statement(ast_module, file_path)
     return functions
 
 
@@ -992,12 +988,21 @@ def filter_files_optimized(file_path: Path, tests_root: Path, ignore_paths: list
 
 def function_has_return_statement(function_node: FunctionDef | AsyncFunctionDef) -> bool:
     # Custom DFS, return True as soon as a Return node is found
-    stack: list[ast.AST] = [function_node]
+    stack: list[ast.AST] = list(function_node.body)
     while stack:
         node = stack.pop()
         if isinstance(node, ast.Return):
             return True
-        stack.extend(ast.iter_child_nodes(node))
+        # Only push child nodes that are statements; Return nodes are statements,
+        # so this preserves correctness while avoiding unnecessary traversal into expr/Name/etc.
+        for field in getattr(node, "_fields", ()):
+            child = getattr(node, field, None)
+            if isinstance(child, list):
+                for item in child:
+                    if isinstance(item, ast.stmt):
+                        stack.append(item)
+            elif isinstance(child, ast.stmt):
+                stack.append(child)
     return False
 
 

codeflash/languages/__init__.py

@@ -19,7 +19,9 @@
 
 from codeflash.languages.base import (
     CodeContext,
+    DependencyResolver,
     HelperFunction,
+    IndexResult,
     Language,
     LanguageSupport,
     ParentInfo,
@@ -82,8 +84,10 @@ def __getattr__(name: str):
 
 __all__ = [
     "CodeContext",
+    "DependencyResolver",
     "FunctionInfo",
     "HelperFunction",
+    "IndexResult",
     "Language",
     "LanguageSupport",
     "ParentInfo",