Merge pull request #1498 from codeflash-ai/cf-simplify-context-extraction

refactor: move context extraction to languages/python/context/

Author: KRRT7

.gitignore

@@ -268,3 +268,5 @@ tessl.json
 
 # Tessl auto-generates AGENTS.md on install; ignore to avoid cluttering git status
 AGENTS.md
+.serena/
+.codeflash/

codeflash/code_utils/config_consts.py

@@ -4,8 +4,8 @@
 from typing import Any, Union
 
 MAX_TEST_RUN_ITERATIONS = 5
-OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 16000
-TESTGEN_CONTEXT_TOKEN_LIMIT = 16000
+OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 48000
+TESTGEN_CONTEXT_TOKEN_LIMIT = 48000
 INDIVIDUAL_TESTCASE_TIMEOUT = 15
 MAX_FUNCTION_TEST_SECONDS = 60
 MIN_IMPROVEMENT_THRESHOLD = 0.05

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.
 

codeflash/languages/current.py

@@ -34,7 +34,7 @@
     from codeflash.languages.base import LanguageSupport
 
 # Module-level singleton for the current language
-_current_language: Language | None = None
+_current_language: Language = Language.PYTHON
 
 
 def current_language() -> Language:

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.
 

codeflash/context/__init__.py …ash/languages/python/context/__init__.pycodeflash/context/__init__.py renamed to codeflash/languages/python/context/__init__.py codeflash/context/__init__.py renamed to codeflash/languages/python/context/__init__.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,69 @@ 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 +488,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 +506,19 @@ 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(