Update code_context_extractor.py

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -142,6 +142,7 @@ def extract_code_string_context_from_files(
     code_context_type: CodeContextType = CodeContextType.READ_ONLY,
 ) -> CodeString:
     """Extract code context from files containing target functions and their helpers.
+
     This function processes two sets of files:
     1. Files containing the function to optimize (fto) and their first-degree helpers
     2. Files containing only helpers of helpers (with no overlap with the first set)
@@ -162,12 +163,12 @@ def extract_code_string_context_from_files(
     """
     # Rearrange to remove overlaps, so we only access each file path once
     helpers_of_helpers_no_overlap = defaultdict(set)
-    for file_path in helpers_of_helpers:
+    for file_path, helper_functions in helpers_of_helpers.items():
         if file_path in helpers_of_fto:
             # Remove duplicates within the same file path, in case a helper of helper is also a helper of fto
             helpers_of_helpers[file_path] -= helpers_of_fto[file_path]
         else:
-            helpers_of_helpers_no_overlap[file_path] = helpers_of_helpers[file_path]
+            helpers_of_helpers_no_overlap[file_path] = helper_functions
 
     final_code_string_context = ""
 
@@ -372,21 +373,19 @@ def get_function_to_optimize_as_function_source(
                 and name.full_name.startswith(name.module_name)
                 and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name
             ):
-                function_source = FunctionSource(
+                return FunctionSource(
                     file_path=function_to_optimize.file_path,
                     qualified_name=function_to_optimize.qualified_name,
                     fully_qualified_name=name.full_name,
                     only_function_name=name.name,
                     source_code=name.get_line_code(),
                     jedi_definition=name,
                 )
-                return function_source
-        except Exception as e:
+        except Exception as e:  # noqa: PERF203
             logger.exception(f"Error while getting function source: {e}")
             continue
-    raise ValueError(
-        f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
-    )
+    msg = f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
+    raise ValueError(msg)
 
 
 def get_function_sources_from_jedi(
@@ -450,13 +449,13 @@ def is_dunder_method(name: str) -> bool:
 
 
 def get_section_names(node: cst.CSTNode) -> list[str]:
-    """Returns the section attribute names (e.g., body, orelse) for a given node if they exist."""
+    """Return the section attribute names (e.g., body, orelse) for a given node if they exist."""
     possible_sections = ["body", "orelse", "finalbody", "handlers"]
     return [sec for sec in possible_sections if hasattr(node, sec)]
 
 
 def remove_docstring_from_body(indented_block: cst.IndentedBlock) -> cst.CSTNode:
-    """Removes the docstring from an indented block if it exists"""
+    """Remove the docstring from an indented block if it exists."""
     if not isinstance(indented_block.body[0], cst.SimpleStatementLine):
         return indented_block
     first_stmt = indented_block.body[0].body[0]
@@ -469,10 +468,12 @@ def parse_code_and_prune_cst(
     code: str,
     code_context_type: CodeContextType,
     target_functions: set[str],
-    helpers_of_helper_functions: set[str] = set(),
-    remove_docstrings: bool = False,
+    helpers_of_helper_functions: set[str] | None = None,
+    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
 ) -> str:
     """Create a read-only version of the code by parsing and filtering the code to keep only class contextual information, and other module scoped variables."""
+    if helpers_of_helper_functions is None:
+        helpers_of_helper_functions = set()
     module = cst.parse_module(code)
     if code_context_type == CodeContextType.READ_WRITABLE:
         filtered_node, found_target = prune_cst_for_read_writable_code(module, target_functions)
@@ -485,7 +486,8 @@ def parse_code_and_prune_cst(
             module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
         )
     else:
-        raise ValueError(f"Unknown code_context_type: {code_context_type}")
+        msg = f"Unknown code_context_type: {code_context_type}"
+        raise TypeError(msg)
 
     if not found_target:
         raise ValueError("No target functions found in the provided code")
@@ -494,7 +496,7 @@ def parse_code_and_prune_cst(
     return ""
 
 
-def prune_cst_for_read_writable_code(
+def prune_cst_for_read_writable_code(  # noqa: PLR0911
     node: cst.CSTNode, target_functions: set[str], prefix: str = ""
 ) -> tuple[cst.CSTNode | None, bool]:
     """Recursively filter the node and its children to build the read-writable codeblock. This contains nodes that lead to target functions.
@@ -520,7 +522,7 @@ def prune_cst_for_read_writable_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            raise TypeError("ClassDef body is not an IndentedBlock")
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
         new_body = []
         found_target = False
@@ -574,14 +576,14 @@ def prune_cst_for_read_writable_code(
     return (node.with_changes(**updates) if updates else node), True
 
 
-def prune_cst_for_read_only_code(
+def prune_cst_for_read_only_code(  # noqa: PLR0911
     node: cst.CSTNode,
     target_functions: set[str],
     helpers_of_helper_functions: set[str],
     prefix: str = "",
-    remove_docstrings: bool = False,
+    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
 ) -> tuple[cst.CSTNode | None, bool]:
-    """Recursively filter the node for read-only context:
+    """Recursively filter the node for read-only context.
 
     Returns:
         (filtered_node, found_target):
@@ -613,7 +615,7 @@ def prune_cst_for_read_only_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            raise TypeError("ClassDef body is not an IndentedBlock")
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
 
@@ -678,14 +680,14 @@ def prune_cst_for_read_only_code(
     return None, False
 
 
-def prune_cst_for_testgen_code(
+def prune_cst_for_testgen_code(  # noqa: PLR0911
     node: cst.CSTNode,
     target_functions: set[str],
     helpers_of_helper_functions: set[str],
     prefix: str = "",
-    remove_docstrings: bool = False,
+    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
 ) -> tuple[cst.CSTNode | None, bool]:
-    """Recursively filter the node for testgen context:
+    """Recursively filter the node for testgen context.
 
     Returns:
         (filtered_node, found_target):
@@ -718,7 +720,7 @@ def prune_cst_for_testgen_code(
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise ValueError("ClassDef body is not an IndentedBlock")
+            raise TypeError("ClassDef body is not an IndentedBlock")
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value