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

KRRT7 · KRRT7 · commit 547c02e8bc48 · 2026-02-16T14:49:04.000-05:00
Move code_context_extractor.py and unused_definition_remover.py from
codeflash/context/ to codeflash/languages/python/context/ and update
all import sites.
diff --git a/codeflash/languages/python/context/__init__.py b/codeflash/languages/python/context/__init__.py
diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py
@@ -14,16 +14,16 @@
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.code_utils.code_utils import encoded_tokens_len, get_qualified_name, path_belongs_to_site_packages
 from codeflash.code_utils.config_consts import OPTIMIZATION_CONTEXT_TOKEN_LIMIT, TESTGEN_CONTEXT_TOKEN_LIMIT
-from codeflash.context.unused_definition_remover import (
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize  # noqa: TC001
+
+# Language support imports for multi-language code context extraction
+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,
     remove_unused_definitions_by_function_names,
 )
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize  # noqa: TC001
-
-# Language support imports for multi-language code context extraction
-from codeflash.languages import Language, is_python
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -38,8 +38,8 @@
 
     from jedi.api.classes import Name
 
-    from codeflash.context.unused_definition_remover import UsageInfo
     from codeflash.languages.base import HelperFunction
+    from codeflash.languages.python.context.unused_definition_remover import UsageInfo
 
 # Error message constants
 READ_WRITABLE_LIMIT_ERROR = "Read-writable code has exceeded token limit, cannot proceed"
diff --git a/codeflash/languages/python/context/unused_definition_remover.py b/codeflash/languages/python/context/unused_definition_remover.py
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -72,15 +72,18 @@
 from codeflash.code_utils.shell_utils import make_env_with_project_root
 from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
 from codeflash.code_utils.time_utils import humanize_runtime
-from codeflash.context import code_context_extractor
-from codeflash.context.unused_definition_remover import detect_unused_helper_functions, revert_unused_helper_functions
 from codeflash.discovery.functions_to_optimize import was_function_previously_optimized
 from codeflash.either import Failure, Success, is_successful
 from codeflash.languages import is_python
 from codeflash.languages.base import Language
 from codeflash.languages.current import current_language_support, is_typescript
 from codeflash.languages.javascript.module_system import detect_module_system
 from codeflash.languages.javascript.test_runner import clear_created_config_files, get_created_config_files
+from codeflash.languages.python.context import code_context_extractor
+from codeflash.languages.python.context.unused_definition_remover import (
+    detect_unused_helper_functions,
+    revert_unused_helper_functions,
+)
 from codeflash.lsp.helpers import is_LSP_enabled, report_to_markdown_table, tree_to_markdown
 from codeflash.lsp.lsp_message import LspCodeMessage, LspMarkdownMessage, LSPMessageId
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
diff --git a/tests/benchmarks/test_benchmark_code_extract_code_context.py b/tests/benchmarks/test_benchmark_code_extract_code_context.py
@@ -1,8 +1,8 @@
 from argparse import Namespace
 from pathlib import Path
 
-from codeflash.context.code_context_extractor import get_code_optimization_context
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 from codeflash.models.models import FunctionParent
 from codeflash.optimization.optimizer import Optimizer
 
diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py
@@ -10,15 +10,15 @@
 
 from codeflash.code_utils.code_extractor import GlobalAssignmentCollector, add_global_assignments
 from codeflash.code_utils.code_replacer import replace_functions_and_add_imports
-from codeflash.context.code_context_extractor import (
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.languages.python.context.code_context_extractor import (
     collect_names_from_annotation,
     enrich_testgen_context,
     extract_classes_from_type_hint,
     extract_imports_for_class,
     get_code_optimization_context,
     resolve_transitive_type_deps,
 )
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import CodeString, CodeStringsMarkdown, FunctionParent
 from codeflash.optimization.optimizer import Optimizer
 
diff --git a/tests/test_get_read_only_code.py b/tests/test_get_read_only_code.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from codeflash.context.code_context_extractor import parse_code_and_prune_cst
+from codeflash.languages.python.context.code_context_extractor import parse_code_and_prune_cst
 from codeflash.models.models import CodeContextType
 
 
diff --git a/tests/test_get_read_writable_code.py b/tests/test_get_read_writable_code.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from codeflash.context.code_context_extractor import parse_code_and_prune_cst
+from codeflash.languages.python.context.code_context_extractor import parse_code_and_prune_cst
 from codeflash.models.models import CodeContextType
 
 
diff --git a/tests/test_get_testgen_code.py b/tests/test_get_testgen_code.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from codeflash.context.code_context_extractor import parse_code_and_prune_cst
+from codeflash.languages.python.context.code_context_extractor import parse_code_and_prune_cst
 from codeflash.models.models import CodeContextType
 
 
diff --git a/tests/test_languages/test_code_context_extraction.py b/tests/test_languages/test_code_context_extraction.py
@@ -20,14 +20,12 @@
 
 from __future__ import annotations
 
-from pathlib import Path
-
 import pytest
 
-from codeflash.context.code_context_extractor import get_code_optimization_context_for_language
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.languages.base import Language
 from codeflash.languages.javascript.support import JavaScriptSupport, TypeScriptSupport
+from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context_for_language
 
 
 @pytest.fixture
diff --git a/tests/test_languages/test_javascript_e2e.py b/tests/test_languages/test_javascript_e2e.py
@@ -106,9 +106,9 @@ def js_project_dir(self):
     def test_extract_code_context_for_javascript(self, js_project_dir):
         """Test extracting code context for a JavaScript function."""
         skip_if_js_not_supported()
-        from codeflash.context.code_context_extractor import get_code_optimization_context
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
         from codeflash.languages import current as lang_current
+        from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 
         lang_current._current_language = Language.JAVASCRIPT
 
diff --git a/tests/test_languages/test_javascript_optimization_flow.py b/tests/test_languages/test_javascript_optimization_flow.py
@@ -9,7 +9,6 @@
 This is the JavaScript equivalent of test_instrument_tests.py for Python.
 """
 
-from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -71,9 +70,9 @@ def test_function_to_optimize_has_correct_language_for_javascript(self, tmp_path
     def test_code_context_preserves_language(self, tmp_path):
         """Verify language is preserved in code context extraction."""
         skip_if_js_not_supported()
-        from codeflash.context.code_context_extractor import get_code_optimization_context
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
         from codeflash.languages import current as lang_current
+        from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 
         lang_current._current_language = Language.TYPESCRIPT
 
@@ -164,7 +163,7 @@ def test_testgen_request_includes_correct_language(self, tmp_path):
 
         # Mock the AI service request
         ai_client = AiServiceClient()
-        with patch.object(ai_client, 'make_ai_service_request') as mock_request:
+        with patch.object(ai_client, "make_ai_service_request") as mock_request:
             mock_response = MagicMock()
             mock_response.status_code = 200
             mock_response.json.return_value = {
@@ -191,8 +190,8 @@ def test_testgen_request_includes_correct_language(self, tmp_path):
             # Verify the request was made with correct language
             assert mock_request.called, "API request should have been made"
             call_args = mock_request.call_args
-            payload = call_args[1].get('payload', call_args[0][1] if len(call_args[0]) > 1 else {})
-            assert payload.get('language') == 'typescript', \
+            payload = call_args[1].get("payload", call_args[0][1] if len(call_args[0]) > 1 else {})
+            assert payload.get("language") == "typescript", \
                 f"Expected language='typescript', got language='{payload.get('language')}'"
 
 
@@ -462,7 +461,7 @@ def test_helper_functions_have_correct_language_javascript(self, tmp_path):
         """Verify helper functions have language='javascript' for .js files."""
         skip_if_js_not_supported()
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
-        from codeflash.languages import current as lang_current, get_language_support
+        from codeflash.languages import current as lang_current
         from codeflash.optimization.function_optimizer import FunctionOptimizer
 
         lang_current._current_language = Language.JAVASCRIPT
diff --git a/tests/test_languages/test_typescript_e2e.py b/tests/test_languages/test_typescript_e2e.py
@@ -69,7 +69,7 @@ def test_discover_functions_with_type_annotations(self):
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
 
         with tempfile.NamedTemporaryFile(suffix=".ts", mode="w", delete=False) as f:
-            f.write("""
+            f.write(r"""
 export function add(a: number, b: number): number {
     return a + b;
 }
@@ -123,9 +123,9 @@ def ts_project_dir(self):
     def test_extract_code_context_for_typescript(self, ts_project_dir):
         """Test extracting code context for a TypeScript function."""
         skip_if_ts_not_supported()
-        from codeflash.context.code_context_extractor import get_code_optimization_context
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
         from codeflash.languages import current as lang_current
+        from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 
         lang_current._current_language = Language.TYPESCRIPT
 
@@ -201,7 +201,7 @@ def test_replace_function_preserves_types(self):
         from codeflash.languages import get_language_support
         from codeflash.languages.base import FunctionInfo
 
-        original_source = """
+        original_source = r"""
 interface Config {
     timeout: number;
     retries: number;
@@ -212,7 +212,7 @@ def test_replace_function_preserves_types(self):
 }
 """
 
-        new_function = """function processConfig(config: Config): string {
+        new_function = r"""function processConfig(config: Config): string {
     // Optimized with template caching
     const { timeout, retries } = config;
     return `timeout=\${timeout}, retries=\${retries}`;
diff --git a/tests/test_languages/test_vitest_e2e.py b/tests/test_languages/test_vitest_e2e.py
@@ -117,10 +117,10 @@ def vitest_project_dir(self):
     def test_extract_code_context_for_typescript(self, vitest_project_dir):
         """Test extracting code context for a TypeScript function."""
         skip_if_js_not_supported()
-        from codeflash.context.code_context_extractor import get_code_optimization_context
         from codeflash.discovery.functions_to_optimize import find_all_functions_in_file
         from codeflash.languages import current as lang_current
         from codeflash.languages.base import Language
+        from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 
         lang_current._current_language = Language.TYPESCRIPT
 
diff --git a/tests/test_remove_unused_definitions.py b/tests/test_remove_unused_definitions.py
@@ -1,6 +1,6 @@
 
 
-from codeflash.context.unused_definition_remover import remove_unused_definitions_by_function_names
+from codeflash.languages.python.context.unused_definition_remover import remove_unused_definitions_by_function_names
 
 
 def test_variable_removal_only() -> None:
diff --git a/tests/test_unused_helper_revert.py b/tests/test_unused_helper_revert.py
@@ -5,8 +5,11 @@
 
 import pytest
 
-from codeflash.context.unused_definition_remover import detect_unused_helper_functions, revert_unused_helper_functions
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.languages.python.context.unused_definition_remover import (
+    detect_unused_helper_functions,
+    revert_unused_helper_functions,
+)
 from codeflash.models.models import CodeStringsMarkdown
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.verification_utils import TestConfig
