Update code_context_extractor.py

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -3,16 +3,19 @@
 import os
 from collections import defaultdict
 from itertools import chain
-from typing import TYPE_CHECKING, Optional
+from pathlib import Path
 
 import jedi
 import libcst as cst
 import tiktoken
+from jedi.api.classes import Name
+from libcst import CSTNode
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
 from codeflash.context.unused_definition_remover import remove_unused_definitions_by_function_names
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -22,14 +25,6 @@
 )
 from codeflash.optimization.function_context import belongs_to_function_qualified
 
-if TYPE_CHECKING:
-    from pathlib import Path
-
-    from jedi.api.classes import Name
-    from libcst import CSTNode
-
-    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-
 
 def get_code_optimization_context(
     function_to_optimize: FunctionToOptimize,
@@ -143,14 +138,13 @@ def extract_code_string_context_from_files(
     helpers_of_fto: dict[Path, set[FunctionSource]],
     helpers_of_helpers: dict[Path, set[FunctionSource]],
     project_root_path: Path,
-    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
+    remove_docstrings: bool = False,
     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).
+    2. Files containing only helpers of helpers (with no overlap with the first set)
 
     For each file, it extracts relevant code based on the specified context type, adds necessary
     imports, and combines them.
@@ -168,12 +162,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, helper_set in helpers_of_helpers.items():
+    for file_path in helpers_of_helpers:
         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_no_overlap[file_path] = helper_set - helpers_of_fto[file_path]
+            helpers_of_helpers[file_path] -= helpers_of_fto[file_path]
         else:
-            helpers_of_helpers_no_overlap[file_path] = helper_set
+            helpers_of_helpers_no_overlap[file_path] = helpers_of_helpers[file_path]
 
     final_code_string_context = ""
 
@@ -250,7 +244,7 @@ def extract_code_markdown_context_from_files(
     helpers_of_fto: dict[Path, set[FunctionSource]],
     helpers_of_helpers: dict[Path, set[FunctionSource]],
     project_root_path: Path,
-    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
+    remove_docstrings: bool = False,
     code_context_type: CodeContextType = CodeContextType.READ_ONLY,
 ) -> CodeStringsMarkdown:
     """Extract code context from files containing target functions and their helpers, formatting them as markdown.
@@ -276,12 +270,12 @@ def extract_code_markdown_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, helper_set in helpers_of_helpers.items():
+    for file_path in helpers_of_helpers:
         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_no_overlap[file_path] = helper_set - helpers_of_fto[file_path]
+            helpers_of_helpers[file_path] -= helpers_of_fto[file_path]
         else:
-            helpers_of_helpers_no_overlap[file_path] = helper_set
+            helpers_of_helpers_no_overlap[file_path] = helpers_of_helpers[file_path]
     code_context_markdown = CodeStringsMarkdown()
     # Extract code from file paths that contain fto and first degree helpers. helpers of helpers may also be included if they are in the same files
     for file_path, function_sources in helpers_of_fto.items():
@@ -378,19 +372,21 @@ 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
             ):
-                return FunctionSource(
+                function_source = 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,
                 )
-        except Exception as e:  # noqa: PERF203
+                return function_source
+        except Exception as e:
             logger.exception(f"Error while getting function source: {e}")
             continue
-    msg = f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
-    raise ValueError(msg)
+    raise ValueError(
+        f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
+    )
 
 
 def get_function_sources_from_jedi(
@@ -411,7 +407,7 @@ def get_function_sources_from_jedi(
             for name in names:
                 try:
                     definitions: list[Name] = name.goto(follow_imports=True, follow_builtin_imports=False)
-                except Exception:
+                except Exception:  # noqa: BLE001
                     logger.debug(f"Error while getting definitions for {qualified_function_name}")
                     definitions = []
                 if definitions:
@@ -454,13 +450,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."""  # noqa: D401
+    """Returns 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:
-    """Remove the docstring from an indented block if it exists."""
+    """Removes 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]
@@ -473,12 +469,10 @@ def parse_code_and_prune_cst(
     code: str,
     code_context_type: CodeContextType,
     target_functions: set[str],
-    helpers_of_helper_functions: Optional[set[str]] = None,
-    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
+    helpers_of_helper_functions: set[str] = set(),
+    remove_docstrings: bool = False,
 ) -> 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)
@@ -491,8 +485,7 @@ def parse_code_and_prune_cst(
             module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
         )
     else:
-        msg = f"Unknown code_context_type: {code_context_type}"
-        raise ValueError(msg)
+        raise ValueError(f"Unknown code_context_type: {code_context_type}")
 
     if not found_target:
         raise ValueError("No target functions found in the provided code")
@@ -501,7 +494,7 @@ def parse_code_and_prune_cst(
     return ""
 
 
-def prune_cst_for_read_writable_code(  # noqa: PLR0911
+def prune_cst_for_read_writable_code(
     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.
@@ -527,7 +520,7 @@ def prune_cst_for_read_writable_code(  # noqa: PLR0911
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise TypeError("ClassDef body is not an IndentedBlock")
+            raise ValueError("ClassDef body is not an IndentedBlock")
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
         new_body = []
         found_target = False
@@ -581,14 +574,14 @@ def prune_cst_for_read_writable_code(  # noqa: PLR0911
     return (node.with_changes(**updates) if updates else node), True
 
 
-def prune_cst_for_read_only_code(  # noqa: PLR0911
+def prune_cst_for_read_only_code(
     node: cst.CSTNode,
     target_functions: set[str],
     helpers_of_helper_functions: set[str],
     prefix: str = "",
-    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
+    remove_docstrings: bool = False,
 ) -> 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):
@@ -620,7 +613,7 @@ def prune_cst_for_read_only_code(  # noqa: PLR0911
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise TypeError("ClassDef body is not an IndentedBlock")
+            raise ValueError("ClassDef body is not an IndentedBlock")
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
 
@@ -685,14 +678,14 @@ def prune_cst_for_read_only_code(  # noqa: PLR0911
     return None, False
 
 
-def prune_cst_for_testgen_code(  # noqa: PLR0911
+def prune_cst_for_testgen_code(
     node: cst.CSTNode,
     target_functions: set[str],
     helpers_of_helper_functions: set[str],
     prefix: str = "",
-    remove_docstrings: bool = False,  # noqa: FBT001, FBT002
+    remove_docstrings: bool = False,
 ) -> tuple[cst.CSTNode | None, bool]:
-    """Recursively filter the node for testgen context.
+    """Recursively filter the node for testgen context:
 
     Returns:
         (filtered_node, found_target):
@@ -725,7 +718,7 @@ def prune_cst_for_testgen_code(  # noqa: PLR0911
             return None, False
         # Assuming always an IndentedBlock
         if not isinstance(node.body, cst.IndentedBlock):
-            raise TypeError("ClassDef body is not an IndentedBlock")
+            raise ValueError("ClassDef body is not an IndentedBlock")
 
         class_prefix = f"{prefix}.{node.name.value}" if prefix else node.name.value
 
@@ -787,4 +780,4 @@ def prune_cst_for_testgen_code(  # noqa: PLR0911
     if updates:
         return (node.with_changes(**updates), found_any_target)
 
-    return None, False
+    return None, False