squash

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -455,39 +455,32 @@ def parse_code_and_prune_cst(
     code: str,
     code_context_type: CodeContextType,
     target_functions: set[str],
-    helpers_of_helper_functions: Optional[set[str]] = None,
+    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)
-
-    dispatch: dict[
-        CodeContextType, Callable[[cst.CSTNode, set[str], set[str], str, bool], tuple[Optional[cst.CSTNode], bool]]
-    ] = {
-        CodeContextType.READ_WRITABLE: prune_cst_for_read_writable_code,
-        CodeContextType.READ_ONLY: prune_cst_for_read_only_code,
-        CodeContextType.TESTGEN: prune_cst_for_testgen_code,
-    }
-
-    prune_func = dispatch.get(code_context_type)
-    if prune_func is None:
+    if code_context_type == CodeContextType.READ_WRITABLE:
+        filtered_node, found_target = prune_cst_for_read_writable_code(module, target_functions)
+    elif code_context_type == CodeContextType.READ_ONLY:
+        filtered_node, found_target = prune_cst_for_read_only_code(
+            module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
+        )
+    elif code_context_type == CodeContextType.TESTGEN:
+        filtered_node, found_target = prune_cst_for_testgen_code(
+            module, target_functions, helpers_of_helper_functions, remove_docstrings=remove_docstrings
+        )
+    else:
         msg = f"Unknown code_context_type: {code_context_type}"
         raise ValueError(msg)
 
-    filtered_node, found_target = prune_func(
-        module, target_functions, helpers_of_helper_functions, "", remove_docstrings
-    )
-
     if not found_target:
         msg = "No target functions found in the provided code"
         raise ValueError(msg)
     if filtered_node and isinstance(filtered_node, cst.Module):
         return str(filtered_node.code)
     return ""
 
-
 def prune_cst_for_read_writable_code(
     node: cst.CSTNode, target_functions: set[str], prefix: str = ""
 ) -> tuple[cst.CSTNode | None, bool]: