Merge remote-tracking branch 'origin/main' into line-profiler

Author: aseembits93

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -1,3 +1,4 @@
+from concurrent.futures import ThreadPoolExecutor
 def funcA(number):
     k = 0
     for i in range(number * 100):
@@ -8,7 +9,14 @@ def funcA(number):
     # Use a generator expression directly in join for more efficiency
     return " ".join(str(i) for i in range(number))
 
+def test_threadpool() -> None:
+    pool = ThreadPoolExecutor(max_workers=3)
+    args = list(range(10, 31, 10))
+    result = pool.map(funcA, args)
+
+    for r in result:
+        print(r)
+
 
 if __name__ == "__main__":
-    for i in range(10, 31, 10):
-        funcA(10)
+    test_threadpool()

codeflash/cli_cmds/cmd_init.py

@@ -595,7 +595,7 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
     if formatter in ["black", "ruff"]:
         try:
             subprocess.run([formatter], capture_output=True, check=False)
-        except FileNotFoundError:
+        except (FileNotFoundError, NotADirectoryError):
             click.echo(f"⚠️ Formatter not found: {formatter}, please ensure it is installed")
     codeflash_section["formatter-cmds"] = formatter_cmds
     # Add the 'codeflash' section, ensuring 'tool' section exists

codeflash/code_utils/code_extractor.py

@@ -235,20 +235,20 @@ def extract_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str |
     return edited_code, contextual_dunder_methods
 
 
-def find_preexisting_objects(source_code: str) -> list[tuple[str, list[FunctionParent]]]:
-    """Find all preexisting functions, classes or class methods in the source code"""
-    preexisting_objects: list[tuple[str, list[FunctionParent]]] = []
+def find_preexisting_objects(source_code: str) -> set[tuple[str, tuple[FunctionParent, ...]]]:
+    """Find all preexisting functions, classes or class methods in the source code."""
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = set()
     try:
         module_node: ast.Module = ast.parse(source_code)
     except SyntaxError:
         logger.exception("find_preexisting_objects - Syntax error while parsing code")
         return preexisting_objects
     for node in module_node.body:
         if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            preexisting_objects.append((node.name, []))
+            preexisting_objects.add((node.name, ()))
         elif isinstance(node, ast.ClassDef):
-            preexisting_objects.append((node.name, []))
+            preexisting_objects.add((node.name, ()))
             for cnode in node.body:
                 if isinstance(cnode, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                    preexisting_objects.append((cnode.name, [FunctionParent(node.name, "ClassDef")]))
+                    preexisting_objects.add((cnode.name, (FunctionParent(node.name, "ClassDef"),)))
     return preexisting_objects

codeflash/code_utils/code_replacer.py

@@ -39,11 +39,11 @@ class OptimFunctionCollector(cst.CSTVisitor):
 
     def __init__(
         self,
-        preexisting_objects: list[tuple[str, list[FunctionParent]]] | None = None,
+        preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] | None = None,
         function_names: set[tuple[str | None, str]] | None = None,
     ) -> None:
         super().__init__()
-        self.preexisting_objects = preexisting_objects if preexisting_objects is not None else []
+        self.preexisting_objects = preexisting_objects if preexisting_objects is not None else set()
 
         self.function_names = function_names  # set of (class_name, function_name)
         self.modified_functions: dict[
@@ -61,7 +61,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
             self.modified_init_functions[self.current_class] = node
         elif (
             self.preexisting_objects
-            and (node.name.value, []) not in self.preexisting_objects
+            and (node.name.value, ()) not in self.preexisting_objects
             and self.current_class is None
         ):
             self.new_functions.append(node)
@@ -72,7 +72,7 @@ def visit_ClassDef(self, node: cst.ClassDef) -> bool:
             return False  # If already in a class, do not recurse deeper
         self.current_class = node.name.value
 
-        parents = [FunctionParent(name=node.name.value, type="ClassDef")]
+        parents = (FunctionParent(name=node.name.value, type="ClassDef"),)
         for child_node in node.body.body:
             if (
                 self.preexisting_objects
@@ -160,7 +160,7 @@ def replace_functions_in_file(
     source_code: str,
     original_function_names: list[str],
     optimized_code: str,
-    preexisting_objects: list[tuple[str, list[FunctionParent]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
 ) -> str:
     parsed_function_names = []
     for original_function_name in original_function_names:
@@ -196,7 +196,7 @@ def replace_functions_and_add_imports(
     function_names: list[str],
     optimized_code: str,
     module_abspath: Path,
-    preexisting_objects: list[tuple[str, list[FunctionParent]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
     project_root_path: Path,
 ) -> str:
     return add_needed_imports_from_module(
@@ -212,7 +212,7 @@ def replace_function_definitions_in_module(
     function_names: list[str],
     optimized_code: str,
     module_abspath: Path,
-    preexisting_objects: list[tuple[str, list[FunctionParent]]],
+    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]],
     project_root_path: Path,
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")

codeflash/context/code_context_extractor.py

@@ -30,18 +30,25 @@ def get_code_optimization_context(
 ) -> CodeOptimizationContext:
     # Get FunctionSource representation of helpers of FTO
     helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi({function_to_optimize.file_path: {function_to_optimize.qualified_name}}, project_root_path)
+
+    # Add function to optimize into helpers of FTO dict, as they'll be processed together
+    fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path)
+    helpers_of_fto_dict[function_to_optimize.file_path].add(fto_as_function_source)
+
+    # Format data to search for helpers of helpers using get_function_sources_from_jedi
     helpers_of_fto_qualified_names_dict = {
         file_path: {source.qualified_name for source in sources}
         for file_path, sources in helpers_of_fto_dict.items()
     }
 
+    # __init__ functions are automatically considered as helpers of FTO, so we add them to the dict (regardless of whether they exist)
+    # This helps us to search for helpers of __init__ functions of classes that contain helpers of FTO
+    for qualified_names in helpers_of_fto_qualified_names_dict.values():
+         qualified_names.update({f"{qn.rsplit('.', 1)[0]}.__init__" for qn in qualified_names if '.' in qn})
+
     # Get FunctionSource representation of helpers of helpers of FTO
     helpers_of_helpers_dict, helpers_of_helpers_list = get_function_sources_from_jedi(helpers_of_fto_qualified_names_dict, project_root_path)
 
-    # Add function to optimize into helpers of FTO dict, as they'll be processed together
-    fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path)
-    helpers_of_fto_dict[function_to_optimize.file_path].add(fto_as_function_source)
-
     # Extract code context for optimization
     final_read_writable_code = extract_code_string_context_from_files(helpers_of_fto_dict,{}, project_root_path, remove_docstrings=False, code_context_type=CodeContextType.READ_WRITABLE).code
     read_only_code_markdown = extract_code_markdown_context_from_files(
@@ -58,8 +65,8 @@ def get_code_optimization_context(
     if final_read_writable_tokens > optim_token_limit:
         raise ValueError("Read-writable code has exceeded token limit, cannot proceed")
 
-    # Setup preexisting objects for code replacer TODO: should remove duplicates
-    preexisting_objects = list(
+    # Setup preexisting objects for code replacer
+    preexisting_objects = set(
         chain(
             find_preexisting_objects(final_read_writable_code),
             *(find_preexisting_objects(codestring.code) for codestring in read_only_code_markdown.code_strings),

codeflash/discovery/discover_unit_tests.py

@@ -186,15 +186,19 @@ def process_test_files(
     jedi_project = jedi.Project(path=project_root_path)
 
     for test_file, functions in file_to_test_map.items():
-        script = jedi.Script(path=test_file, project=jedi_project)
-        test_functions = set()
-
-        all_names = script.get_names(all_scopes=True, references=True)
-        all_defs = script.get_names(all_scopes=True, definitions=True)
-        all_names_top = script.get_names(all_scopes=True)
-
-        top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
-        top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}
+        try:
+            script = jedi.Script(path=test_file, project=jedi_project)
+            test_functions = set()
+
+            all_names = script.get_names(all_scopes=True, references=True)
+            all_defs = script.get_names(all_scopes=True, definitions=True)
+            all_names_top = script.get_names(all_scopes=True)
+
+            top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
+            top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}
+        except Exception as e:
+            logger.debug(f"Failed to get jedi script for {test_file}: {e}")
+            continue
 
         if test_framework == "pytest":
             for function in functions:

codeflash/models/models.py

@@ -99,7 +99,7 @@ class CodeOptimizationContext(BaseModel):
     read_writable_code: str = Field(min_length=1)
     read_only_context_code: str = ""
     helper_functions: list[FunctionSource]
-    preexisting_objects: list[tuple[str, list[FunctionParent]]]
+    preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]]
 
 class CodeContextType(str, Enum):
     READ_WRITABLE = "READ_WRITABLE"