todo cleanup

Author: aseembits93

codeflash/code_utils/code_extractor.py

@@ -34,3 +34,6 @@
 N_CANDIDATES_LP_EFFECTIVE = min(N_CANDIDATES_LP_LSP if _IS_LSP_ENABLED else N_CANDIDATES_LP, MAX_N_CANDIDATES_LP)
 N_TESTS_TO_GENERATE_EFFECTIVE = N_TESTS_TO_GENERATE_LSP if _IS_LSP_ENABLED else N_TESTS_TO_GENERATE
 TOTAL_LOOPING_TIME_EFFECTIVE = TOTAL_LOOPING_TIME_LSP if _IS_LSP_ENABLED else TOTAL_LOOPING_TIME
+
+MAX_CONTEXT_LEN_IMPACT = 1000
+TIME_LIMIT_FOR_OPT_IMPACT = 5  # in sec

codeflash/code_utils/config_consts.py

@@ -377,7 +377,6 @@ def generate_and_instrument_tests(
 
     # note: this isn't called by the lsp, only called by cli
     def optimize_function(self) -> Result[BestOptimization, str]:
-        get_opt_impact_metrics(self.function_to_optimize.file_path,self.function_to_optimize.qualified_name, self.project_root, self.test_cfg.tests_root)
         initialization_result = self.can_be_optimized()
         if not is_successful(initialization_result):
             return Failure(initialization_result.failure())
@@ -1470,7 +1469,11 @@ def process_review(
                 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
+                self.function_to_optimize_source_code,
+                self.function_to_optimize.file_path,
+                self.function_to_optimize.qualified_name,
+                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

codeflash/optimization/function_optimizer.py

@@ -1,28 +1,131 @@
 #!/usr/bin/env python3
-"""Example of using the ripgrep search script programmatically."""
+"""Example usage of the function reference finder."""
 
-import json
+from find_function_references import find_function_references
+from find_function_references_detailed import find_function_references_detailed
 
-from ripgrep_search import search_with_ripgrep
 
-# Search for any pattern you want
-pattern = "sorter"  # Change this to any pattern you need
-results = search_with_ripgrep(pattern)
+def example_basic_usage():
+    """Example of basic usage - just get list of files."""
+    # Example: Find references to a function
+    filepath = "path/to/your/file.py"
+    function_name = "your_function_name"
 
-# Access the results as a dictionary
-print(f"Found matches in {len(results)} files")
+    # Find references (will auto-detect project root)
+    reference_files = find_function_references(filepath, function_name)
 
-# 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]}...")
+    print(f"Files containing references to {function_name}:")
+    for file in reference_files:
+        print(f"  - {file}")
 
-# Save results to a JSON file if needed
-with open("search_results.json", "w") as f:
-    json.dump(results, f, indent=2)
+    return reference_files
 
-# 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)}")
+def example_detailed_usage():
+    """Example of detailed usage - get line numbers and context."""
+    # Example: Find references with detailed information
+    filepath = "path/to/your/file.py"
+    function_name = "your_function_name"
+    project_root = "/path/to/project/root"  # Optional
+
+    # Find detailed references
+    references = find_function_references_detailed(filepath, function_name, project_root)
+
+    print(f"\nDetailed references to {function_name}:")
+    for file, refs in references.items():
+        print(f"\nFile: {file}")
+        for ref in refs:
+            print(f"  Line {ref['line']}: {ref['context']}")
+
+    return references
+
+
+def example_programmatic_usage():
+    """Example of using in your own Python code."""
+    import jedi
+
+    # Direct Jedi usage for more control
+    source_code = """
+def my_function(x, y):
+    return x + y
+
+result = my_function(1, 2)
+"""
+
+    # Create a script object
+    script = jedi.Script(code=source_code)
+
+    # Find the function definition (line 2, column 4 for 'my_function')
+    references = script.get_references(line=2, column=4)
+
+    print("Direct Jedi references:")
+    for ref in references:
+        print(f"  Line {ref.line}, Column {ref.column}: {ref.description}")
+
+    # You can also search for names
+    names = script.get_names()
+    for name in names:
+        if name.type == "function":
+            print(f"Found function: {name.name} at line {name.line}")
+
+
+def example_find_all_functions_and_their_references():
+    """Example of finding all functions in a file and their references."""
+    import os
+
+    import jedi
+
+    def find_all_functions_with_references(filepath: str):
+        """Find all functions in a file and their references."""
+        with open(filepath) as f:
+            source_code = f.read()
+
+        script = jedi.Script(code=source_code, path=filepath)
+
+        # Get all defined names
+        names = script.get_names()
+
+        functions_and_refs = {}
+
+        for name in names:
+            if name.type == "function":
+                # Get references for this function
+                refs = script.get_references(line=name.line, column=name.column)
+
+                ref_locations = []
+                for ref in refs:
+                    if ref.module_path and str(ref.module_path) != filepath:
+                        ref_locations.append({"file": str(ref.module_path), "line": ref.line, "column": ref.column})
+
+                functions_and_refs[name.name] = {"definition_line": name.line, "references": ref_locations}
+
+        return functions_and_refs
+
+    # Example usage
+    filepath = "your_module.py"
+    if os.path.exists(filepath):
+        all_refs = find_all_functions_with_references(filepath)
+
+        for func_name, info in all_refs.items():
+            print(f"\nFunction: {func_name} (defined at line {info['definition_line']})")
+            if info["references"]:
+                print("  Referenced in:")
+                for ref in info["references"]:
+                    print(f"    - {ref['file']}:{ref['line']}")
+            else:
+                print("  No external references found")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Function Reference Finder - Usage Examples")
+    print("=" * 60)
+
+    print("\nNote: Update the file paths and function names in the examples")
+    print("before running them with your actual code.\n")
+
+    # Uncomment to run examples:
+    # example_basic_usage()
+    # example_detailed_usage()
+    # example_programmatic_usage()
+    # example_find_all_functions_and_their_references()

example_usage.py

@@ -1,26 +1,25 @@
 #!/usr/bin/env python3
-"""
-Script to find all references to the sorter function from code_to_optimize/bubble_sort.py
+"""Script to find all references to the sorter function from code_to_optimize/bubble_sort.py
 using Jedi's static analysis capabilities.
 """
 
-import jedi
-import os
 from pathlib import Path
 
+import jedi
+
 
 def find_function_references(file_path, line, column, project_root):
-    """
-    Find all references to a function using Jedi.
+    """Find all references to a function using Jedi.
 
     Args:
         file_path: Path to the file containing the function
         line: Line number where the function is defined (1-indexed)
         column: Column number where the function name starts (0-indexed)
         project_root: Root directory of the project to search
+
     """
     # Read the source code
-    with open(file_path, 'r') as f:
+    with open(file_path) as f:
         source = f.read()
 
     # Create a Jedi Script object with project configuration
@@ -45,7 +44,7 @@ def find_function_references(file_path, line, column, project_root):
     references = []
     try:
         # Use usages() method to get all references
-        references = script.get_references(line, column, scope='project', include_builtins=False)
+        references = script.get_references(line, column, scope="project", include_builtins=False)
     except AttributeError:
         # Alternative approach using search
         print("Using alternative search method...")
@@ -91,7 +90,7 @@ def main():
             for ref in sorted(file_refs, key=lambda r: (r.line, r.column)):
                 # Get the line content for context
                 try:
-                    with open(file_path, 'r') as f:
+                    with open(file_path) as f:
                         lines = f.readlines()
                         if ref.line <= len(lines):
                             line_content = lines[ref.line - 1].strip()
@@ -108,4 +107,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

find_sorter_references.py

@@ -143,8 +143,8 @@ def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
 
                 # For methods, add proper indentation (4 spaces)
                 if self.current_class_stack:
-                    lines = func_source.split('\n')
-                    func_source = '\n'.join('    ' + line if line else line for line in lines)
+                    lines = func_source.split("\n")
+                    func_source = "\n".join("    " + line if line else line for line in lines)
 
                 self.function_definitions[full_name] = FunctionDefinitionInfo(
                     name=full_name,
@@ -153,7 +153,7 @@ def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
                     start_line=position.start.line if position else -1,
                     end_line=position.end.line if position else -1,
                     is_method=bool(self.current_class_stack),
-                    class_name=self.current_class_stack[-1] if self.current_class_stack else None
+                    class_name=self.current_class_stack[-1] if self.current_class_stack else None,
                 )
 
             # Handle nested functions - mark parent as containing nested calls
@@ -167,11 +167,15 @@ def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
                     parent_source = cst.Module(body=[parent_func_node]).code
 
                     # Get parent class context (go up one level in stack since we're inside the nested function)
-                    parent_class_stack = self.current_class_stack[:-1] if len(self.current_function_stack) == 1 and self.current_class_stack else []
+                    parent_class_stack = (
+                        self.current_class_stack[:-1]
+                        if len(self.current_function_stack) == 1 and self.current_class_stack
+                        else []
+                    )
 
                     if parent_class_stack:
-                        lines = parent_source.split('\n')
-                        parent_source = '\n'.join('    ' + line if line else line for line in lines)
+                        lines = parent_source.split("\n")
+                        parent_source = "\n".join("    " + line if line else line for line in lines)
 
                     self.function_definitions[parent_name] = FunctionDefinitionInfo(
                         name=parent_name,
@@ -180,7 +184,7 @@ def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
                         start_line=parent_position.start.line if parent_position else -1,
                         end_line=parent_position.end.line if parent_position else -1,
                         is_method=bool(parent_class_stack),
-                        class_name=parent_class_stack[-1] if parent_class_stack else None
+                        class_name=parent_class_stack[-1] if parent_class_stack else None,
                     )
 
             # Reset the flag for parent function if we're in nested functions
@@ -290,10 +294,7 @@ def get_results(self) -> Dict[str, str]:
             Only includes functions that call the target function (directly or through nested functions).
 
         """
-        return {
-            info.name: info.source_code
-            for info in self.function_definitions.values()
-        }
+        return {info.name: info.source_code for info in self.function_definitions.values()}
 
 
 def find_function_calls(source_code: str, target_function_name: str, target_filepath: str) -> Dict[str, str]:
@@ -367,11 +368,12 @@ def inner():
 
     # Simple usage - results is just a dict of {function_name: source_code}
     import json
+
     print("JSON representation of results:")
     print(json.dumps(list(results.keys()), indent=2))
 
     print("\nFormatted output:")
     for func_name, source_code in results.items():
         print(f"\n=== {func_name} ===")
         print(source_code)
-        print()
+        print()