refactor: deduplicate Python language support code

KRRT7 · KRRT7 · commit 707703ca59cc · 2026-02-16T16:55:01.000-05:00
Extract shared helpers and remove dead code across the language support area:
- Extract `is_assignment_used()` and move `recurse_sections` to unused_definition_remover.py, replacing duplicated logic in both context files
- Extract `function_sources_to_helpers()` in support.py to unify identical HelperFunction construction
- Remove dead `get_comment_prefix()` method from protocol and all implementations (comment_prefix property serves all callers)
diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py
@@ -519,15 +519,6 @@ def get_test_file_suffix(self) -> str:
         """
         ...
 
-    def get_comment_prefix(self) -> str:
-        """Get the comment prefix for this language.
-
-        Returns:
-            Comment prefix (e.g., "//" for JS, "#" for Python).
-
-        """
-        ...
-
     def find_test_root(self, project_root: Path) -> Path | None:
         """Find the test root directory for a project.
 
diff --git a/codeflash/languages/javascript/support.py b/codeflash/languages/javascript/support.py
@@ -1805,15 +1805,6 @@ def get_test_file_suffix(self) -> str:
         """
         return ".test.js"
 
-    def get_comment_prefix(self) -> str:
-        """Get the comment prefix for JavaScript.
-
-        Returns:
-            JavaScript single-line comment prefix.
-
-        """
-        return "//"
-
     def find_test_root(self, project_root: Path) -> Path | None:
         """Find the test root directory for a JavaScript project.
 
diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py
@@ -20,8 +20,9 @@
 from codeflash.languages import Language, is_python
 from codeflash.languages.python.context.unused_definition_remover import (
     collect_top_level_defs_with_usages,
-    extract_names_from_targets,
     get_section_names,
+    is_assignment_used,
+    recurse_sections,
     remove_unused_definitions_by_function_names,
 )
 from codeflash.models.models import (
@@ -34,8 +35,6 @@
 from codeflash.optimization.function_context import belongs_to_function_qualified
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
-
     from jedi.api.classes import Name
 
     from codeflash.languages.base import HelperFunction
@@ -1103,50 +1102,6 @@ def _validate_classdef(node: cst.ClassDef, prefix: str) -> tuple[str, cst.Indent
     return _qualified_name(prefix, node.name.value), node.body
 
 
-def _recurse_sections(
-    node: cst.CSTNode,
-    section_names: list[str],
-    prune_fn: Callable[[cst.CSTNode], tuple[cst.CSTNode | None, bool]],
-    keep_non_target_children: bool = False,
-) -> tuple[cst.CSTNode | None, bool]:
-    updates: dict[str, list[cst.CSTNode] | cst.CSTNode] = {}
-    found_any_target = False
-    for section in section_names:
-        original_content = getattr(node, section, None)
-        if isinstance(original_content, (list, tuple)):
-            new_children = []
-            section_found_target = False
-            for child in original_content:
-                filtered, found_target = prune_fn(child)
-                if filtered:
-                    new_children.append(filtered)
-                section_found_target |= found_target
-            if keep_non_target_children:
-                if section_found_target or new_children:
-                    found_any_target |= section_found_target
-                    updates[section] = new_children
-            elif section_found_target:
-                found_any_target = True
-                updates[section] = new_children
-        elif original_content is not None:
-            filtered, found_target = prune_fn(original_content)
-            if keep_non_target_children:
-                found_any_target |= found_target
-                if filtered:
-                    updates[section] = filtered
-            elif found_target:
-                found_any_target = True
-                if filtered:
-                    updates[section] = filtered
-    if keep_non_target_children:
-        if updates:
-            return node.with_changes(**updates), found_any_target
-        return None, False
-    if not found_any_target:
-        return None, False
-    return (node.with_changes(**updates) if updates else node), True
-
-
 def prune_cst(
     node: cst.CSTNode,
     target_functions: set[str],
@@ -1278,19 +1233,9 @@ def prune_cst(
 
     # Handle assignments for READ_WRITABLE mode
     if defs_with_usages is not None:
-        if isinstance(node, cst.Assign):
-            for target in node.targets:
-                names = extract_names_from_targets(target.target)
-                for name in names:
-                    if name in defs_with_usages and defs_with_usages[name].used_by_qualified_function:
-                        return node, True
-            return None, False
-
-        if isinstance(node, (cst.AnnAssign, cst.AugAssign)):
-            names = extract_names_from_targets(node.target)
-            for name in names:
-                if name in defs_with_usages and defs_with_usages[name].used_by_qualified_function:
-                    return node, True
+        if isinstance(node, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
+            if is_assignment_used(node, defs_with_usages):
+                return node, True
             return None, False
 
     # For other nodes, recursively process children
@@ -1299,7 +1244,7 @@ def prune_cst(
         return node, False
 
     if helpers is not None:
-        return _recurse_sections(
+        return recurse_sections(
             node,
             section_names,
             lambda child: prune_cst(
@@ -1317,7 +1262,7 @@ def prune_cst(
             ),
             keep_non_target_children=True,
         )
-    return _recurse_sections(
+    return recurse_sections(
         node,
         section_names,
         lambda child: prune_cst(
diff --git a/codeflash/languages/python/context/unused_definition_remover.py b/codeflash/languages/python/context/unused_definition_remover.py
@@ -15,6 +15,8 @@
 from codeflash.models.models import CodeString, CodeStringsMarkdown
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.models.models import CodeOptimizationContext, FunctionSource
 
@@ -49,6 +51,73 @@ def extract_names_from_targets(target: cst.CSTNode) -> list[str]:
     return names
 
 
+def is_assignment_used(
+    node: cst.CSTNode,
+    definitions: dict[str, UsageInfo],
+    name_prefix: str = "",
+) -> bool:
+    if isinstance(node, cst.Assign):
+        for target in node.targets:
+            names = extract_names_from_targets(target.target)
+            for name in names:
+                lookup = f"{name_prefix}{name}" if name_prefix else name
+                if lookup in definitions and definitions[lookup].used_by_qualified_function:
+                    return True
+        return False
+    if isinstance(node, (cst.AnnAssign, cst.AugAssign)):
+        names = extract_names_from_targets(node.target)
+        for name in names:
+            lookup = f"{name_prefix}{name}" if name_prefix else name
+            if lookup in definitions and definitions[lookup].used_by_qualified_function:
+                return True
+        return False
+    return False
+
+
+def recurse_sections(
+    node: cst.CSTNode,
+    section_names: list[str],
+    prune_fn: Callable[[cst.CSTNode], tuple[cst.CSTNode | None, bool]],
+    keep_non_target_children: bool = False,
+) -> tuple[cst.CSTNode | None, bool]:
+    updates: dict[str, list[cst.CSTNode] | cst.CSTNode] = {}
+    found_any_target = False
+    for section in section_names:
+        original_content = getattr(node, section, None)
+        if isinstance(original_content, (list, tuple)):
+            new_children = []
+            section_found_target = False
+            for child in original_content:
+                filtered, found_target = prune_fn(child)
+                if filtered:
+                    new_children.append(filtered)
+                section_found_target |= found_target
+            if keep_non_target_children:
+                if section_found_target or new_children:
+                    found_any_target |= section_found_target
+                    updates[section] = new_children
+            elif section_found_target:
+                found_any_target = True
+                updates[section] = new_children
+        elif original_content is not None:
+            filtered, found_target = prune_fn(original_content)
+            if keep_non_target_children:
+                found_any_target |= found_target
+                if filtered:
+                    updates[section] = filtered
+            elif found_target:
+                found_any_target = True
+                if filtered:
+                    updates[section] = filtered
+    if keep_non_target_children:
+        if updates:
+            return node.with_changes(**updates), found_any_target
+        return None, False
+    if not found_any_target:
+        return None, False
+    return (node.with_changes(**updates) if updates else node), True
+
+
 def collect_top_level_definitions(
     node: cst.CSTNode, definitions: Optional[dict[str, UsageInfo]] = None
 ) -> dict[str, UsageInfo]:
@@ -423,27 +492,9 @@ def remove_unused_definitions_recursively(
                 elif isinstance(statement, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
                     var_used = False
 
-                    # Check if any variable in this assignment is used
-                    if isinstance(statement, cst.Assign):
-                        for target in statement.targets:
-                            names = extract_names_from_targets(target.target)
-                            for name in names:
-                                class_var_name = f"{class_name}.{name}"
-                                if (
-                                    class_var_name in definitions
-                                    and definitions[class_var_name].used_by_qualified_function
-                                ):
-                                    var_used = True
-                                    method_or_var_used = True
-                                    break
-                    elif isinstance(statement, (cst.AnnAssign, cst.AugAssign)):
-                        names = extract_names_from_targets(statement.target)
-                        for name in names:
-                            class_var_name = f"{class_name}.{name}"
-                            if class_var_name in definitions and definitions[class_var_name].used_by_qualified_function:
-                                var_used = True
-                                method_or_var_used = True
-                                break
+                    if is_assignment_used(statement, definitions, name_prefix=f"{class_name}."):
+                        var_used = True
+                        method_or_var_used = True
 
                     if var_used or class_has_dependencies:
                         new_statements.append(statement)
@@ -459,56 +510,21 @@ def remove_unused_definitions_recursively(
 
         return node, method_or_var_used or class_has_dependencies
 
-    # Handle assignments (Assign and AnnAssign)
-    if isinstance(node, cst.Assign):
-        for target in node.targets:
-            names = extract_names_from_targets(target.target)
-            for name in names:
-                if name in definitions and definitions[name].used_by_qualified_function:
-                    return node, True
-        return None, False
-
-    if isinstance(node, (cst.AnnAssign, cst.AugAssign)):
-        names = extract_names_from_targets(node.target)
-        for name in names:
-            if name in definitions and definitions[name].used_by_qualified_function:
-                return node, True
+    # Handle assignments (Assign, AnnAssign, AugAssign)
+    if isinstance(node, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
+        if is_assignment_used(node, definitions):
+            return node, True
         return None, False
 
     # For other nodes, recursively process children
     section_names = get_section_names(node)
     if not section_names:
         return node, False
-
-    updates = {}
-    found_used = False
-
-    for section in section_names:
-        original_content = getattr(node, section, None)
-        if isinstance(original_content, (list, tuple)):
-            new_children = []
-            section_found_used = False
-
-            for child in original_content:
-                filtered, used = remove_unused_definitions_recursively(child, definitions)
-                if filtered:
-                    new_children.append(filtered)
-                section_found_used |= used
-
-            if new_children or section_found_used:
-                found_used |= section_found_used
-                updates[section] = new_children
-        elif original_content is not None:
-            filtered, used = remove_unused_definitions_recursively(original_content, definitions)
-            found_used |= used
-            if filtered:
-                updates[section] = filtered
-    if not found_used:
-        return None, False
-    if updates:
-        return node.with_changes(**updates), found_used
-
-    return node, False
+    return recurse_sections(
+        node,
+        section_names,
+        lambda child: remove_unused_definitions_recursively(child, definitions),
+    )
 
 
 def collect_top_level_defs_with_usages(
diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py
