Update code_context_extractor.py

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -375,9 +375,7 @@ def get_function_to_optimize_as_function_source(
             )
 
     msg = f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}"
-    raise ValueError(
-        msg
-    )
+    raise ValueError(msg)
 
 
 def get_function_sources_from_jedi(
@@ -463,20 +461,22 @@ def parse_code_and_prune_cst(
     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)
-    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:
+
+    dispatch = {
+        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:
         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=remove_docstrings
+    )
+
     if not found_target:
         msg = "No target functions found in the provided code"
         raise ValueError(msg)
@@ -625,10 +625,12 @@ def prune_cst_for_read_only_code(
             return None, False
 
         if remove_docstrings:
-            return node.with_changes(
-                body=remove_docstring_from_body(node.body.with_changes(body=new_class_body))
-            ) if new_class_body else None, True
-        return node.with_changes(body=node.body.with_changes(body=new_class_body)) if new_class_body else None, True
+            return (
+                node.with_changes(body=remove_docstring_from_body(node.body.with_changes(body=new_class_body)))
+                if new_class_body
+                else None
+            ), True
+        return (node.with_changes(body=node.body.with_changes(body=new_class_body)) if new_class_body else None), True
 
     # For other nodes, keep the node and recursively filter children
     section_names = get_section_names(node)
@@ -731,10 +733,12 @@ def prune_cst_for_testgen_code(
             return None, False
 
         if remove_docstrings:
-            return node.with_changes(
-                body=remove_docstring_from_body(node.body.with_changes(body=new_class_body))
-            ) if new_class_body else None, True
-        return node.with_changes(body=node.body.with_changes(body=new_class_body)) if new_class_body else None, True
+            return (
+                node.with_changes(body=remove_docstring_from_body(node.body.with_changes(body=new_class_body)))
+                if new_class_body
+                else None
+            ), True
+        return (node.with_changes(body=node.body.with_changes(body=new_class_body)) if new_class_body else None), True
 
     # For other nodes, keep the node and recursively filter children
     section_names = get_section_names(node)