wip

Author: aseembits93

codeflash/code_utils/code_extractor.py

@@ -2,7 +2,10 @@
 from __future__ import annotations
 
 import ast
+import json
+import subprocess
 from itertools import chain
+from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
@@ -14,12 +17,10 @@
 from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
-    from pathlib import Path
-
     from libcst.helpers import ModuleNameAndPackage
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-    from codeflash.models.models import FunctionSource
+    from codeflash.models.models import FunctionSource, ImpactMetrics
 
 
 class GlobalAssignmentCollector(cst.CSTVisitor):
@@ -748,3 +749,85 @@ def find_preexisting_objects(source_code: str) -> set[tuple[str, tuple[FunctionP
                 if isinstance(cnode, (ast.FunctionDef, ast.AsyncFunctionDef)):
                     preexisting_objects.add((cnode.name, (FunctionParent(node.name, "ClassDef"),)))
     return preexisting_objects
+
+
+def search_with_ripgrep(pattern: str, path: str = ".") -> dict[str, list[tuple[int, str]]]:
+    """Use ripgrep to search for a pattern in the repository.
+
+    Args:
+        pattern: The pattern to search for
+        path: The directory to search in (default: current directory)
+
+    Returns:
+        Dictionary with filepaths as keys and list of (line_no, content) tuples as values
+
+    """
+    # Run ripgrep with JSON output for easier parsing
+    # -n: Show line numbers
+    # --json: Output in JSON format
+    # --no-heading: Don't group matches by file
+    path = str(Path.cwd())
+    cmd = [
+        "rg",
+        "-n",
+        "--json",
+        pattern,
+        path,
+        "-g",
+        "!/Users/aseemsaxena/Downloads/codeflash_dev/codeflash/code_to_optimize/tests/**",
+    ]
+    print(" ".join(cmd))
+    # Parse the JSON output
+    matches_dict = {}
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False,  # Don't raise exception on non-zero return
+        )
+
+        if result.returncode not in [0, 1]:  # 0 = matches found, 1 = no matches
+            print(f"Error running ripgrep: {result.stderr}")
+            return {}
+
+        for line in result.stdout.strip().split("\n"):
+            if not line:
+                continue
+
+            try:
+                json_obj = json.loads(line)
+
+                # We're only interested in match objects
+                if json_obj.get("type") == "match":
+                    data = json_obj.get("data", {})
+                    file_path = data.get("path", {}).get("text", "")
+                    line_number = data.get("line_number")
+                    line_content = data.get("lines", {}).get("text", "").rstrip("\n")
+
+                    if file_path and line_number:
+                        if file_path not in matches_dict:
+                            matches_dict[file_path] = []
+                        matches_dict[file_path].append((line_number, line_content))
+
+            except json.JSONDecodeError:
+                continue
+
+    except FileNotFoundError:
+        print("Error: ripgrep (rg) is not installed or not in PATH")
+        return {}
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+        return {}
+    return matches_dict
+
+
+def get_opt_impact_metrics(file_path: Path, qualified_name: str, project_root: Path, tests_root: Path) -> ImpactMetrics:
+    # grep for function / use rg (respects gitignore)
+    # SAFE_GREP_EXECUTABLE command
+    # ast visitor for occurances and loop occurances
+    # radon lib for complexity metrics
+    print(file_path, qualified_name, project_root, tests_root)
+
+    # grep windows alternative
+    return 0

codeflash/code_utils/compat.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import sys
 import tempfile
 from pathlib import Path
@@ -17,6 +18,7 @@ class Compat:
     LF: str = os.linesep
 
     SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix()
+    SAFE_GREP_EXECUTABLE: str = shutil.which("grep")  # works even grep is aliased in the env
 
     IS_POSIX: bool = os.name != "nt"
 
@@ -45,3 +47,4 @@ def codeflash_cache_db(self) -> Path:
 LF = _compat.LF
 SAFE_SYS_EXECUTABLE = _compat.SAFE_SYS_EXECUTABLE
 IS_POSIX = _compat.IS_POSIX
+SAFE_GREP_EXECUTABLE = _compat.SAFE_GREP_EXECUTABLE

codeflash/models/models.py

@@ -31,6 +31,14 @@
 from codeflash.verification.comparator import comparator
 
 
+@dataclass
+class ImpactMetrics:
+    complexity_score: int
+    occurances: int
+    loop_occurances: int
+    presence_of_decorators: bool
+
+
 @dataclass(frozen=True)
 class AIServiceRefinerRequest:
     optimization_id: str

codeflash/optimization/function_optimizer.py

@@ -24,6 +24,7 @@
 from codeflash.benchmarking.utils import process_benchmark_data
 from codeflash.cli_cmds.console import code_print, console, logger, lsp_log, progress_bar
 from codeflash.code_utils import env_utils
+from codeflash.code_utils.code_extractor import get_opt_impact_metrics
 from codeflash.code_utils.code_replacer import (
     add_custom_marker_to_all_tests,
     modify_autouse_fixture,
@@ -1467,6 +1468,9 @@ def process_review(
             except Exception as e:
                 logger.debug(f"optimization impact response failed, investigate {e}")
             data["optimization_impact"] = opt_impact_response
+            data["impact_metrics"] = get_opt_impact_metrics(
+                self.project_root, self.test_cfg.tests_root
+            )  # need module root, tests root only
         if raise_pr and not staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)

example_usage.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+"""
+Example of using the ripgrep search script programmatically.
+"""
+
+from ripgrep_search import search_with_ripgrep
+import json
+
+# Search for any pattern you want
+pattern = "sorter"  # Change this to any pattern you need
+results = search_with_ripgrep(pattern)
+
+# Access the results as a dictionary
+print(f"Found matches in {len(results)} files")
+
+# Iterate through the results
+for filepath, occurrences in results.items():
+    print(f"\n{filepath}: {len(occurrences)} matches")
+    for line_no, content in occurrences[:3]:  # Show first 3 matches per file
+        print(f"  Line {line_no}: {content[:80]}...")
+
+# Save results to a JSON file if needed
+with open("search_results.json", "w") as f:
+    json.dump(results, f, indent=2)
+
+# Or filter results for specific files
+python_files_only = {
+    path: matches
+    for path, matches in results.items()
+    if path.endswith('.py')
+}
+
+print(f"\nPython files with matches: {len(python_files_only)}")