From b59882e1fbf70b602da1046f83d93776dd32033a Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Wed, 4 Feb 2026 16:48:15 +0000 Subject: [PATCH 1/5] prevent pytest fail with non unicode chars --- codeflash/code_utils/shell_utils.py | 4 ++++ codeflash/languages/javascript/find_references.py | 4 ++-- codeflash/languages/javascript/parse.py | 7 +------ codeflash/languages/javascript/vitest_runner.py | 3 +-- codeflash/languages/treesitter_utils.py | 4 ++-- codeflash/optimization/function_optimizer.py | 6 +++--- codeflash/verification/parse_test_output.py | 8 +++----- 7 files changed, 16 insertions(+), 20 deletions(-) diff --git a/codeflash/code_utils/shell_utils.py b/codeflash/code_utils/shell_utils.py index df2cff2d6..61dc9f2fd 100644 --- a/codeflash/code_utils/shell_utils.py +++ b/codeflash/code_utils/shell_utils.py @@ -247,6 +247,10 @@ def get_cross_platform_subprocess_run_args( capture_output: bool = True, ) -> dict[str, str]: run_args = {"cwd": cwd, "env": env, "text": text, "timeout": timeout, "check": check} + # When text=True, use errors='replace' to handle non-UTF-8 bytes gracefully + # instead of raising UnicodeDecodeError + if text: + run_args["errors"] = "replace" if sys.platform == "win32": creationflags = subprocess.CREATE_NEW_PROCESS_GROUP run_args["creationflags"] = creationflags diff --git a/codeflash/languages/javascript/find_references.py b/codeflash/languages/javascript/find_references.py index 16b93cfca..3b48761dd 100644 --- a/codeflash/languages/javascript/find_references.py +++ b/codeflash/languages/javascript/find_references.py @@ -213,7 +213,7 @@ def find_references( trigger_check = True if import_info: context.visited_files.add(file_path) - import_name, original_import = import_info + import_name, original_import = import_info # noqa: RUF059 file_refs = self._find_references_in_file( file_path, file_code, reexport_name, import_name, file_analyzer, include_self=True ) @@ -404,7 +404,7 @@ def _find_identifier_references( name_node = node.child_by_field_name("name") if name_node: new_current_function = source_bytes[name_node.start_byte : name_node.end_byte].decode("utf8") - elif node.type in ("variable_declarator",): + elif node.type in ("variable_declarator",): # noqa: FURB171 # Arrow function or function expression assigned to variable name_node = node.child_by_field_name("name") value_node = node.child_by_field_name("value") diff --git a/codeflash/languages/javascript/parse.py b/codeflash/languages/javascript/parse.py index bc24013d9..d6b43feae 100644 --- a/codeflash/languages/javascript/parse.py +++ b/codeflash/languages/javascript/parse.py @@ -16,12 +16,7 @@ from junitparser.xunit2 import JUnitXml from codeflash.cli_cmds.console import logger -from codeflash.models.models import ( - FunctionTestInvocation, - InvocationId, - TestResults, - TestType, -) +from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType if TYPE_CHECKING: import subprocess diff --git a/codeflash/languages/javascript/vitest_runner.py b/codeflash/languages/javascript/vitest_runner.py index a5f6552d3..f622d7384 100644 --- a/codeflash/languages/javascript/vitest_runner.py +++ b/codeflash/languages/javascript/vitest_runner.py @@ -288,8 +288,7 @@ def run_vitest_behavioral_tests( logger.debug(f"Vitest JUnit XML created: {result_file_path} ({file_size} bytes)") if file_size < 200: # Suspiciously small - likely empty or just headers logger.warning( - f"Vitest JUnit XML is very small ({file_size} bytes). " - f"Content: {result_file_path.read_text()[:500]}" + f"Vitest JUnit XML is very small ({file_size} bytes). Content: {result_file_path.read_text()[:500]}" ) else: logger.warning( diff --git a/codeflash/languages/treesitter_utils.py b/codeflash/languages/treesitter_utils.py index f4b7ead43..8125161c1 100644 --- a/codeflash/languages/treesitter_utils.py +++ b/codeflash/languages/treesitter_utils.py @@ -1580,9 +1580,9 @@ def get_analyzer_for_file(file_path: Path) -> TreeSitterAnalyzer: """ suffix = file_path.suffix.lower() - if suffix in (".ts",): + if suffix in (".ts",): # noqa: FURB171 return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT) - if suffix in (".tsx",): + if suffix in (".tsx",): # noqa: FURB171 return TreeSitterAnalyzer(TreeSitterLanguage.TSX) # Default to JavaScript for .js, .jsx, .mjs, .cjs return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index ad39557c1..89b19d02c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -315,7 +315,7 @@ def _handle_empty_queue(self) -> CandidateNode | None: self.future_all_code_repair, "Repairing {0} candidates", "Added {0} candidates from repair, total candidates now: {1}", - lambda: self.future_all_code_repair.clear(), + lambda: self.future_all_code_repair.clear(), # noqa: PLW0108 ) if self.line_profiler_done and not self.refinement_done: return self._process_candidates( @@ -330,7 +330,7 @@ def _handle_empty_queue(self) -> CandidateNode | None: self.future_adaptive_optimizations, "Applying adaptive optimizations to {0} candidates", "Added {0} candidates from adaptive optimization, total candidates now: {1}", - lambda: self.future_adaptive_optimizations.clear(), + lambda: self.future_adaptive_optimizations.clear(), # noqa: PLW0108 ) return None # All done @@ -2093,7 +2093,7 @@ def process_review( formatted_generated_test = format_generated_code(concolic_test_str, self.args.formatter_cmds) generated_tests_str += f"```{code_lang}\n{formatted_generated_test}\n```\n\n" - existing_tests, replay_tests, concolic_tests = existing_tests_source_for( + existing_tests, replay_tests, concolic_tests = existing_tests_source_for( # noqa: RUF059 self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root), function_to_all_tests, test_cfg=self.test_cfg, diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py index c80a287e5..4c2c809eb 100644 --- a/codeflash/verification/parse_test_output.py +++ b/codeflash/verification/parse_test_output.py @@ -1,6 +1,5 @@ from __future__ import annotations -import contextlib import os import re import sqlite3 @@ -22,6 +21,9 @@ ) from codeflash.discovery.discover_unit_tests import discover_parameters_unittest from codeflash.languages import is_javascript + +# Import Jest-specific parsing from the JavaScript language module +from codeflash.languages.javascript.parse import parse_jest_test_xml as _parse_jest_test_xml from codeflash.models.models import ( ConcurrencyMetrics, FunctionTestInvocation, @@ -32,10 +34,6 @@ ) from codeflash.verification.coverage_utils import CoverageUtils, JestCoverageUtils -# Import Jest-specific parsing from the JavaScript language module -from codeflash.languages.javascript.parse import jest_end_pattern, jest_start_pattern -from codeflash.languages.javascript.parse import parse_jest_test_xml as _parse_jest_test_xml - if TYPE_CHECKING: import subprocess From 8be10e24ba7e8ca54e7f68791ce3cd2728d36f77 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Wed, 4 Feb 2026 17:01:46 +0000 Subject: [PATCH 2/5] fix: re-export jest patterns for backward compatibility --- codeflash/verification/parse_test_output.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py index 4c2c809eb..bd6511334 100644 --- a/codeflash/verification/parse_test_output.py +++ b/codeflash/verification/parse_test_output.py @@ -23,7 +23,11 @@ from codeflash.languages import is_javascript # Import Jest-specific parsing from the JavaScript language module -from codeflash.languages.javascript.parse import parse_jest_test_xml as _parse_jest_test_xml +from codeflash.languages.javascript.parse import ( + jest_end_pattern, + jest_start_pattern, + parse_jest_test_xml as _parse_jest_test_xml, +) from codeflash.models.models import ( ConcurrencyMetrics, FunctionTestInvocation, From 1a0bb05bb85179c6063b7de48c0eb06a80ce75a2 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 4 Feb 2026 09:19:47 -0800 Subject: [PATCH 3/5] undo non relevant changes --- codeflash/verification/parse_test_output.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py index bd6511334..c80a287e5 100644 --- a/codeflash/verification/parse_test_output.py +++ b/codeflash/verification/parse_test_output.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextlib import os import re import sqlite3 @@ -21,13 +22,6 @@ ) from codeflash.discovery.discover_unit_tests import discover_parameters_unittest from codeflash.languages import is_javascript - -# Import Jest-specific parsing from the JavaScript language module -from codeflash.languages.javascript.parse import ( - jest_end_pattern, - jest_start_pattern, - parse_jest_test_xml as _parse_jest_test_xml, -) from codeflash.models.models import ( ConcurrencyMetrics, FunctionTestInvocation, @@ -38,6 +32,10 @@ ) from codeflash.verification.coverage_utils import CoverageUtils, JestCoverageUtils +# Import Jest-specific parsing from the JavaScript language module +from codeflash.languages.javascript.parse import jest_end_pattern, jest_start_pattern +from codeflash.languages.javascript.parse import parse_jest_test_xml as _parse_jest_test_xml + if TYPE_CHECKING: import subprocess From ff4c36599ba33e0bba7a30946c0f9dc491629c37 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 4 Feb 2026 10:13:11 -0800 Subject: [PATCH 4/5] mypy fix --- codeflash/github/PrComment.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 7416329bb..7ea94ba93 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -31,7 +31,7 @@ def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[B if name: report_table[name] = result - result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = { + result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = { # type: ignore[no-redef] "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -45,10 +45,10 @@ def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[B } if self.original_async_throughput is not None and self.best_async_throughput is not None: - result["original_async_throughput"] = str(self.original_async_throughput) - result["best_async_throughput"] = str(self.best_async_throughput) + result["original_async_throughput"] = str(self.original_async_throughput) # type: ignore[assignment] + result["best_async_throughput"] = str(self.best_async_throughput) # type: ignore[assignment] - return result + return result # type: ignore[return-value] class FileDiffContent(BaseModel): From a1483b5b5aa01d6584e1f32500db4c3daf2cf97b Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 25 Feb 2026 22:36:23 +0530 Subject: [PATCH 5/5] only necessary changes --- codeflash/github/PrComment.py | 12 +- codeflash/languages/__init__.py | 4 + codeflash/languages/base.py | 149 +- codeflash/languages/current.py | 2 +- .../languages/javascript/code_replacer.py | 217 ++ codeflash/languages/javascript/edit_tests.py | 185 +- .../languages/javascript/find_references.py | 12 +- .../languages/javascript/import_resolver.py | 13 +- codeflash/languages/javascript/instrument.py | 440 +++- .../languages/javascript/line_profiler.py | 2 +- .../languages/javascript/module_system.py | 35 +- codeflash/languages/javascript/parse.py | 135 +- codeflash/languages/javascript/support.py | 225 +- codeflash/languages/javascript/test_runner.py | 266 ++- codeflash/languages/javascript/treesitter.py | 1826 +++++++++++++++++ .../languages/javascript/vitest_runner.py | 301 ++- codeflash/languages/python/__init__.py | 3 +- .../languages/python/context/__init__.py | 0 .../python/context/code_context_extractor.py | 1508 ++++++++++++++ .../context/unused_definition_remover.py | 854 ++++++++ codeflash/languages/python/reference_graph.py | 544 +++++ .../python/static_analysis/__init__.py | 0 .../python/static_analysis/code_extractor.py | 1737 ++++++++++++++++ .../python/static_analysis/code_replacer.py | 689 +++++++ .../python/static_analysis/concolic_utils.py | 126 ++ .../python/static_analysis/coverage_utils.py | 93 + .../static_analysis/edit_generated_tests.py | 272 +++ .../static_analysis/line_profile_utils.py | 388 ++++ .../python/static_analysis/static_analysis.py | 167 ++ codeflash/languages/python/support.py | 239 ++- 30 files changed, 10166 insertions(+), 278 deletions(-) create mode 100644 codeflash/languages/javascript/code_replacer.py create mode 100644 codeflash/languages/javascript/treesitter.py create mode 100644 codeflash/languages/python/context/__init__.py create mode 100644 codeflash/languages/python/context/code_context_extractor.py create mode 100644 codeflash/languages/python/context/unused_definition_remover.py create mode 100644 codeflash/languages/python/reference_graph.py create mode 100644 codeflash/languages/python/static_analysis/__init__.py create mode 100644 codeflash/languages/python/static_analysis/code_extractor.py create mode 100644 codeflash/languages/python/static_analysis/code_replacer.py create mode 100644 codeflash/languages/python/static_analysis/concolic_utils.py create mode 100644 codeflash/languages/python/static_analysis/coverage_utils.py create mode 100644 codeflash/languages/python/static_analysis/edit_generated_tests.py create mode 100644 codeflash/languages/python/static_analysis/line_profile_utils.py create mode 100644 codeflash/languages/python/static_analysis/static_analysis.py diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 7ea94ba93..85c24ec57 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -26,12 +26,12 @@ class PrComment: def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]]: report_table: dict[str, dict[str, int]] = {} - for test_type, result in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items(): + for test_type, counts in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items(): name = test_type.to_name() if name: - report_table[name] = result + report_table[name] = counts - result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = { # type: ignore[no-redef] + result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = { "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -45,10 +45,10 @@ def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[B } if self.original_async_throughput is not None and self.best_async_throughput is not None: - result["original_async_throughput"] = str(self.original_async_throughput) # type: ignore[assignment] - result["best_async_throughput"] = str(self.best_async_throughput) # type: ignore[assignment] + result["original_async_throughput"] = self.original_async_throughput + result["best_async_throughput"] = self.best_async_throughput - return result # type: ignore[return-value] + return result class FileDiffContent(BaseModel): diff --git a/codeflash/languages/__init__.py b/codeflash/languages/__init__.py index 47136f4e7..daf33b43c 100644 --- a/codeflash/languages/__init__.py +++ b/codeflash/languages/__init__.py @@ -19,7 +19,9 @@ from codeflash.languages.base import ( CodeContext, + DependencyResolver, HelperFunction, + IndexResult, Language, LanguageSupport, ParentInfo, @@ -82,8 +84,10 @@ def __getattr__(name: str): __all__ = [ "CodeContext", + "DependencyResolver", "FunctionInfo", "HelperFunction", + "IndexResult", "Language", "LanguageSupport", "ParentInfo", diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 99cefdf46..3e10da319 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -11,10 +11,11 @@ from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Iterable, Sequence from pathlib import Path from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.models.models import FunctionSource, GeneratedTestsList, InvocationId from codeflash.languages.language_enum import Language from codeflash.models.function_types import FunctionParent @@ -34,6 +35,16 @@ def __getattr__(name: str) -> Any: raise AttributeError(msg) +@dataclass(frozen=True) +class IndexResult: + file_path: Path + cached: bool + num_edges: int + edges: tuple[tuple[str, str, bool], ...] # (caller_qn, callee_name, is_cross_file) + cross_file_edges: int + error: bool + + @dataclass class HelperFunction: """A helper function that is a dependency of the target function. @@ -192,6 +203,35 @@ class ReferenceInfo: caller_function: str | None = None +@runtime_checkable +class DependencyResolver(Protocol): + """Protocol for language-specific dependency resolution. + + Implementations analyze source files to discover call-graph edges + between functions so the optimizer can extract richer context. + """ + + def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: + """Pre-index a batch of files.""" + ... + + def get_callees( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: + """Return callees for the given functions.""" + ... + + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[tuple[Path, str], int]: + """Return the number of callees for each (file_path, qualified_name) pair.""" + ... + + def close(self) -> None: + """Release resources (e.g. database connections).""" + ... + + @runtime_checkable class LanguageSupport(Protocol): """Protocol defining what a language implementation must provide. @@ -254,6 +294,14 @@ def comment_prefix(self) -> str: """Like # or //.""" ... + @property + def dir_excludes(self) -> frozenset[str]: + """Directory name patterns to skip during file discovery. + + Supports glob wildcards: "name" for exact, "prefix*" for startswith, "*suffix" for endswith. + """ + ... + # === Discovery === def discover_functions( @@ -490,6 +538,87 @@ def remove_test_functions(self, test_source: str, functions_to_remove: list[str] """ ... + def postprocess_generated_tests( + self, generated_tests: GeneratedTestsList, test_framework: str, project_root: Path, source_file_path: Path + ) -> GeneratedTestsList: + """Apply language-specific postprocessing to generated tests. + + Args: + generated_tests: Generated tests to update. + test_framework: Test framework used for the project. + project_root: Project root directory. + source_file_path: Path to the source file under optimization. + + Returns: + Updated generated tests. + + """ + ... + + def remove_test_functions_from_generated_tests( + self, generated_tests: GeneratedTestsList, functions_to_remove: list[str] + ) -> GeneratedTestsList: + """Remove specific test functions from generated tests. + + Args: + generated_tests: Generated tests to update. + functions_to_remove: List of function names to remove. + + Returns: + Updated generated tests. + + """ + ... + + def add_runtime_comments_to_generated_tests( + self, + generated_tests: GeneratedTestsList, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], + tests_project_rootdir: Path | None = None, + ) -> GeneratedTestsList: + """Add runtime comments to generated tests. + + Args: + generated_tests: Generated tests to update. + original_runtimes: Mapping of invocation IDs to original runtimes. + optimized_runtimes: Mapping of invocation IDs to optimized runtimes. + tests_project_rootdir: Root directory for tests (if applicable). + + Returns: + Updated generated tests. + + """ + ... + + def add_global_declarations(self, optimized_code: str, original_source: str, module_abspath: Path) -> str: + """Add new global declarations from optimized code to original source. + + Args: + optimized_code: The optimized code that may contain new declarations. + original_source: The original source code. + module_abspath: Path to the module file (for parser selection). + + Returns: + Original source with new declarations added. + + """ + ... + + def extract_calling_function_source(self, source_code: str, function_name: str, ref_line: int) -> str | None: + """Extract the source code of a calling function. + + Args: + source_code: Full source code of the file. + function_name: Name of the function to extract. + ref_line: Line number where the reference is. + + Returns: + Source code of the function, or None if not found. + + """ + ... + # === Test Result Comparison === def compare_test_results( @@ -519,15 +648,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. @@ -565,6 +685,15 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: # Default implementation: just copy runtime files return False + def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | None: + """Create a language-specific dependency resolver, if available. + + Returns: + A DependencyResolver instance, or None if not supported. + + """ + return None + def instrument_existing_test( self, test_path: Path, diff --git a/codeflash/languages/current.py b/codeflash/languages/current.py index ecdb7315a..005249669 100644 --- a/codeflash/languages/current.py +++ b/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: diff --git a/codeflash/languages/javascript/code_replacer.py b/codeflash/languages/javascript/code_replacer.py new file mode 100644 index 000000000..83c96ec6a --- /dev/null +++ b/codeflash/languages/javascript/code_replacer.py @@ -0,0 +1,217 @@ +"""JavaScript/TypeScript code replacement helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from codeflash.cli_cmds.console import logger + +if TYPE_CHECKING: + from pathlib import Path + + from codeflash.languages.base import Language + from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer + + +# Author: ali +def _add_global_declarations_for_language( + optimized_code: str, original_source: str, module_abspath: Path, language: Language +) -> str: + """Add new global declarations from optimized code to original source. + + Finds module-level declarations (const, let, var, class, type, interface, enum) + in the optimized code that don't exist in the original source and adds them. + + New declarations are inserted after any existing declarations they depend on. + For example, if optimized code has `const _has = FOO.bar.bind(FOO)`, and `FOO` + is already declared in the original source, `_has` will be inserted after `FOO`. + + Args: + optimized_code: The optimized code that may contain new declarations. + original_source: The original source code. + module_abspath: Path to the module file (for parser selection). + language: The language of the code. + + Returns: + Original source with new declarations added in dependency order. + + """ + from codeflash.languages.base import Language + + if language not in (Language.JAVASCRIPT, Language.TYPESCRIPT): + return original_source + + try: + from codeflash.languages.javascript.treesitter import get_analyzer_for_file + + analyzer = get_analyzer_for_file(module_abspath) + + original_declarations = analyzer.find_module_level_declarations(original_source) + optimized_declarations = analyzer.find_module_level_declarations(optimized_code) + + if not optimized_declarations: + return original_source + + existing_names = _get_existing_names(original_declarations, analyzer, original_source) + new_declarations = _filter_new_declarations(optimized_declarations, existing_names) + + if not new_declarations: + return original_source + + # Build a map of existing declaration names to their end lines (1-indexed) + existing_decl_end_lines = {decl.name: decl.end_line for decl in original_declarations} + + # Insert each new declaration after its dependencies + result = original_source + for decl in new_declarations: + result = _insert_declaration_after_dependencies( + result, decl, existing_decl_end_lines, analyzer, module_abspath + ) + # Update the map with the newly inserted declaration for subsequent insertions + # Re-parse to get accurate line numbers after insertion + updated_declarations = analyzer.find_module_level_declarations(result) + existing_decl_end_lines = {d.name: d.end_line for d in updated_declarations} + + return result + + except Exception as e: + logger.debug(f"Error adding global declarations: {e}") + return original_source + + +# Author: ali +def _get_existing_names(original_declarations: list, analyzer: TreeSitterAnalyzer, original_source: str) -> set[str]: + """Get all names that already exist in the original source (declarations + imports).""" + existing_names = {decl.name for decl in original_declarations} + + original_imports = analyzer.find_imports(original_source) + for imp in original_imports: + if imp.default_import: + existing_names.add(imp.default_import) + for name, alias in imp.named_imports: + existing_names.add(alias if alias else name) + if imp.namespace_import: + existing_names.add(imp.namespace_import) + + return existing_names + + +# Author: ali +def _filter_new_declarations(optimized_declarations: list, existing_names: set[str]) -> list: + """Filter declarations to only those that don't exist in the original source.""" + new_declarations = [] + seen_sources: set[str] = set() + + # Sort by line number to maintain order from optimized code + sorted_declarations = sorted(optimized_declarations, key=lambda d: d.start_line) + + for decl in sorted_declarations: + if decl.name not in existing_names and decl.source_code not in seen_sources: + new_declarations.append(decl) + seen_sources.add(decl.source_code) + + return new_declarations + + +# Author: ali +def _insert_declaration_after_dependencies( + source: str, + declaration, + existing_decl_end_lines: dict[str, int], + analyzer: TreeSitterAnalyzer, + module_abspath: Path, +) -> str: + """Insert a declaration after the last existing declaration it depends on. + + Args: + source: Current source code. + declaration: The declaration to insert. + existing_decl_end_lines: Map of existing declaration names to their end lines. + analyzer: TreeSitter analyzer. + module_abspath: Path to the module file. + + Returns: + Source code with the declaration inserted at the correct position. + + """ + # Find identifiers referenced in this declaration + referenced_names = analyzer.find_referenced_identifiers(declaration.source_code) + + # Find the latest end line among all referenced declarations + insertion_line = _find_insertion_line_for_declaration(source, referenced_names, existing_decl_end_lines, analyzer) + + lines = source.splitlines(keepends=True) + + # Ensure proper spacing + decl_code = declaration.source_code + if not decl_code.endswith("\n"): + decl_code += "\n" + + # Add blank line before if inserting after content + if insertion_line > 0 and lines[insertion_line - 1].strip(): + decl_code = "\n" + decl_code + + before = lines[:insertion_line] + after = lines[insertion_line:] + + return "".join([*before, decl_code, *after]) + + +# Author: ali +def _find_insertion_line_for_declaration( + source: str, referenced_names: set[str], existing_decl_end_lines: dict[str, int], analyzer: TreeSitterAnalyzer +) -> int: + """Find the line where a declaration should be inserted based on its dependencies. + + Args: + source: Source code. + referenced_names: Names referenced by the declaration. + existing_decl_end_lines: Map of declaration names to their end lines (1-indexed). + analyzer: TreeSitter analyzer. + + Returns: + Line index (0-based) where the declaration should be inserted. + + """ + # Find the maximum end line among referenced declarations + max_dependency_line = 0 + for name in referenced_names: + if name in existing_decl_end_lines: + max_dependency_line = max(max_dependency_line, existing_decl_end_lines[name]) + + if max_dependency_line > 0: + # Insert after the last dependency (end_line is 1-indexed, we need 0-indexed) + return max_dependency_line + + # No dependencies found - insert after imports + lines = source.splitlines(keepends=True) + return _find_line_after_imports(lines, analyzer, source) + + +# Author: ali +def _find_line_after_imports(lines: list[str], analyzer: TreeSitterAnalyzer, source: str) -> int: + """Find the line index after all imports. + + Args: + lines: Source lines. + analyzer: TreeSitter analyzer. + source: Full source code. + + Returns: + Line index (0-based) for insertion after imports. + + """ + try: + imports = analyzer.find_imports(source) + if imports: + return max(imp.end_line for imp in imports) + except Exception as exc: + logger.debug(f"Exception in _find_line_after_imports: {exc}") + + # Default: insert at beginning (after shebang/directive comments) + for i, line in enumerate(lines): + stripped = line.strip() + if stripped and not stripped.startswith("//") and not stripped.startswith("#!"): + return i + + return 0 diff --git a/codeflash/languages/javascript/edit_tests.py b/codeflash/languages/javascript/edit_tests.py index a4523e83b..601da3cda 100644 --- a/codeflash/languages/javascript/edit_tests.py +++ b/codeflash/languages/javascript/edit_tests.py @@ -6,29 +6,13 @@ from __future__ import annotations +import os import re +from pathlib import Path from codeflash.cli_cmds.console import logger -from codeflash.code_utils.time_utils import format_perf, format_time -from codeflash.result.critic import performance_gain - - -def format_runtime_comment(original_time: int, optimized_time: int) -> str: - """Format a runtime comparison comment for JavaScript. - - Args: - original_time: Original runtime in nanoseconds. - optimized_time: Optimized runtime in nanoseconds. - - Returns: - Formatted comment string with // prefix. - - """ - perf_gain = format_perf( - abs(performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) * 100) - ) - status = "slower" if optimized_time > original_time else "faster" - return f"// {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})" +from codeflash.code_utils.time_utils import format_runtime_comment +from codeflash.models.models import GeneratedTests, GeneratedTestsList def add_runtime_comments(source: str, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int]) -> str: @@ -117,7 +101,7 @@ def find_matching_test(test_description: str) -> str | None: # Only add comment if line has a function call and doesn't already have a comment if func_call_pattern.search(line) and "//" not in line and "expect(" in line: orig_time, opt_time = timing_by_full_name[current_matched_full_name] - comment = format_runtime_comment(orig_time, opt_time) + comment = format_runtime_comment(orig_time, opt_time, comment_prefix="//") logger.debug(f"[js-annotations] Adding comment to test '{current_test_name}': {comment}") # Add comment at end of line line = f"{line.rstrip()} {comment}" @@ -130,6 +114,165 @@ def find_matching_test(test_description: str) -> str | None: return "\n".join(modified_lines) +JS_TEST_EXTENSIONS = ( + ".test.ts", + ".test.js", + ".test.tsx", + ".test.jsx", + ".spec.ts", + ".spec.js", + ".spec.tsx", + ".spec.jsx", + ".ts", + ".js", + ".tsx", + ".jsx", + ".mjs", + ".mts", +) + + +# TODO:{self} Needs cleanup for jest logic in else block +# Author: Sarthak Agarwal +def is_js_test_module_path(test_module_path: str) -> bool: + """Return True when the module path looks like a JS/TS test path.""" + return any(test_module_path.endswith(ext) for ext in JS_TEST_EXTENSIONS) + + +# Author: Sarthak Agarwal +def resolve_js_test_module_path(test_module_path: str, tests_project_rootdir: Path) -> Path: + """Resolve a JS/TS test module path to a concrete file path.""" + if "/" in test_module_path or "\\" in test_module_path: + return tests_project_rootdir / Path(test_module_path) + + matched_ext = None + for ext in JS_TEST_EXTENSIONS: + if test_module_path.endswith(ext): + matched_ext = ext + break + + if matched_ext: + base_path = test_module_path[: -len(matched_ext)] + file_path = base_path.replace(".", os.sep) + matched_ext + tests_dir_name = tests_project_rootdir.name + if file_path.startswith((tests_dir_name + os.sep, tests_dir_name + "/")): + return tests_project_rootdir.parent / Path(file_path) + return tests_project_rootdir / Path(file_path) + + return tests_project_rootdir / Path(test_module_path) + + +# Patterns for normalizing codeflash imports (legacy -> npm package) +# Author: Sarthak Agarwal +_CODEFLASH_REQUIRE_PATTERN = re.compile( + r"(const|let|var)\s+(\w+)\s*=\s*require\s*\(\s*['\"]\.?/?codeflash-jest-helper['\"]\s*\)" +) +_CODEFLASH_IMPORT_PATTERN = re.compile(r"import\s+(?:\*\s+as\s+)?(\w+)\s+from\s+['\"]\.?/?codeflash-jest-helper['\"]") + + +# Author: Sarthak Agarwal +def normalize_codeflash_imports(source: str) -> str: + """Normalize codeflash imports to use the npm package. + + Replaces legacy local file imports: + const codeflash = require('./codeflash-jest-helper') + import codeflash from './codeflash-jest-helper' + + With npm package imports: + const codeflash = require('codeflash') + + Args: + source: JavaScript/TypeScript source code. + + Returns: + Source code with normalized imports. + + """ + # Replace CommonJS require + source = _CODEFLASH_REQUIRE_PATTERN.sub(r"\1 \2 = require('codeflash')", source) + # Replace ES module import + return _CODEFLASH_IMPORT_PATTERN.sub(r"import \1 from 'codeflash'", source) + + +# Author: ali +def inject_test_globals(generated_tests: GeneratedTestsList, test_framework: str = "jest") -> GeneratedTestsList: + # TODO: inside the prompt tell the llm if it should import jest functions or it's already injected in the global window + """Inject test globals into all generated tests. + + Args: + generated_tests: List of generated tests. + test_framework: The test framework being used ("jest", "vitest", or "mocha"). + + Returns: + Generated tests with test globals injected. + + """ + # we only inject test globals for esm modules + # Use vitest imports for vitest projects, jest imports for jest projects + if test_framework == "vitest": + global_import = "import { vi, describe, it, expect, beforeEach, afterEach, beforeAll, test } from 'vitest'\n" + else: + # Default to jest imports for jest and other frameworks + global_import = ( + "import { jest, describe, it, expect, beforeEach, afterEach, beforeAll, test } from '@jest/globals'\n" + ) + + for test in generated_tests.generated_tests: + test.generated_original_test_source = global_import + test.generated_original_test_source + test.instrumented_behavior_test_source = global_import + test.instrumented_behavior_test_source + test.instrumented_perf_test_source = global_import + test.instrumented_perf_test_source + return generated_tests + + +# Author: ali +def disable_ts_check(generated_tests: GeneratedTestsList) -> GeneratedTestsList: + """Disable TypeScript type checking in all generated tests. + + Args: + generated_tests: List of generated tests. + + Returns: + Generated tests with TypeScript type checking disabled. + + """ + # we only inject test globals for esm modules + ts_nocheck = "// @ts-nocheck\n" + + for test in generated_tests.generated_tests: + test.generated_original_test_source = ts_nocheck + test.generated_original_test_source + test.instrumented_behavior_test_source = ts_nocheck + test.instrumented_behavior_test_source + test.instrumented_perf_test_source = ts_nocheck + test.instrumented_perf_test_source + return generated_tests + + +# Author: Sarthak Agarwal +def normalize_generated_tests_imports(generated_tests: GeneratedTestsList) -> GeneratedTestsList: + """Normalize codeflash imports in all generated tests. + + Args: + generated_tests: List of generated tests. + + Returns: + Generated tests with normalized imports. + + """ + normalized_tests = [] + for test in generated_tests.generated_tests: + # Only normalize JS/TS files + if test.behavior_file_path.suffix in (".js", ".ts", ".jsx", ".tsx", ".mjs", ".mts"): + normalized_test = GeneratedTests( + generated_original_test_source=normalize_codeflash_imports(test.generated_original_test_source), + instrumented_behavior_test_source=normalize_codeflash_imports(test.instrumented_behavior_test_source), + instrumented_perf_test_source=normalize_codeflash_imports(test.instrumented_perf_test_source), + behavior_file_path=test.behavior_file_path, + perf_file_path=test.perf_file_path, + ) + normalized_tests.append(normalized_test) + else: + normalized_tests.append(test) + return GeneratedTestsList(generated_tests=normalized_tests) + + def remove_test_functions(source: str, functions_to_remove: list[str]) -> str: """Remove specific test functions from JavaScript test source code. diff --git a/codeflash/languages/javascript/find_references.py b/codeflash/languages/javascript/find_references.py index 3b48761dd..ed6e30636 100644 --- a/codeflash/languages/javascript/find_references.py +++ b/codeflash/languages/javascript/find_references.py @@ -23,7 +23,7 @@ from tree_sitter import Node from codeflash.discovery.functions_to_optimize import FunctionToOptimize - from codeflash.languages.treesitter_utils import ImportInfo, TreeSitterAnalyzer + from codeflash.languages.javascript.treesitter import ImportInfo, TreeSitterAnalyzer logger = logging.getLogger(__name__) @@ -112,7 +112,7 @@ def find_references( List of Reference objects describing each call site. """ - from codeflash.languages.treesitter_utils import get_analyzer_for_file + from codeflash.languages.javascript.treesitter import get_analyzer_for_file function_name = function_to_optimize.function_name source_file = function_to_optimize.file_path @@ -168,7 +168,7 @@ def find_references( if import_info: # Found an import - mark as visited and search for calls context.visited_files.add(file_path) - import_name, original_import = import_info + import_name, _original_import = import_info file_refs = self._find_references_in_file( file_path, file_code, function_name, import_name, file_analyzer, include_self=True ) @@ -213,7 +213,7 @@ def find_references( trigger_check = True if import_info: context.visited_files.add(file_path) - import_name, original_import = import_info # noqa: RUF059 + import_name, _original_import = import_info file_refs = self._find_references_in_file( file_path, file_code, reexport_name, import_name, file_analyzer, include_self=True ) @@ -404,7 +404,7 @@ def _find_identifier_references( name_node = node.child_by_field_name("name") if name_node: new_current_function = source_bytes[name_node.start_byte : name_node.end_byte].decode("utf8") - elif node.type in ("variable_declarator",): # noqa: FURB171 + elif node.type == "variable_declarator": # Arrow function or function expression assigned to variable name_node = node.child_by_field_name("name") value_node = node.child_by_field_name("value") @@ -719,7 +719,7 @@ def _find_reexports_direct( continue # Create a fake ImportInfo to resolve the re-export source - from codeflash.languages.treesitter_utils import ImportInfo + from codeflash.languages.javascript.treesitter import ImportInfo fake_import = ImportInfo( module_path=exp.reexport_source, diff --git a/codeflash/languages/javascript/import_resolver.py b/codeflash/languages/javascript/import_resolver.py index 4e237b8d6..b5ec67115 100644 --- a/codeflash/languages/javascript/import_resolver.py +++ b/codeflash/languages/javascript/import_resolver.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import HelperFunction - from codeflash.languages.treesitter_utils import ImportInfo, TreeSitterAnalyzer + from codeflash.languages.javascript.treesitter import ImportInfo, TreeSitterAnalyzer logger = logging.getLogger(__name__) @@ -44,8 +44,7 @@ def __init__(self, project_root: Path) -> None: project_root: Root directory of the project. """ - # Resolve to real path to handle macOS symlinks like /var -> /private/var - self.project_root = project_root.resolve() + self.project_root = project_root self._resolution_cache: dict[tuple[Path, str], Path | None] = {} def resolve_import(self, import_info: ImportInfo, source_file: Path) -> ResolvedImport | None: @@ -486,7 +485,7 @@ def _extract_helper_from_file( """ from codeflash.languages.base import HelperFunction - from codeflash.languages.treesitter_utils import get_analyzer_for_file + from codeflash.languages.javascript.treesitter import get_analyzer_for_file try: source = file_path.read_text(encoding="utf-8") @@ -558,7 +557,8 @@ def _find_helpers_recursive( """ from codeflash.discovery.functions_to_optimize import FunctionToOptimize - from codeflash.languages.treesitter_utils import get_analyzer_for_file + from codeflash.languages.javascript.treesitter import get_analyzer_for_file + from codeflash.languages.registry import get_language_support if context.current_depth >= context.max_depth: return {} @@ -578,12 +578,15 @@ def _find_helpers_recursive( imports = analyzer.find_imports(source) # Create FunctionToOptimize for the helper + # Get language from the language support registry + lang_support = get_language_support(file_path) func_info = FunctionToOptimize( function_name=helper.name, file_path=file_path, parents=[], starting_line=helper.start_line, ending_line=helper.end_line, + language=str(lang_support.language), ) # Recursively find helpers diff --git a/codeflash/languages/javascript/instrument.py b/codeflash/languages/javascript/instrument.py index 30e7fff7a..8bcd0b2ee 100644 --- a/codeflash/languages/javascript/instrument.py +++ b/codeflash/languages/javascript/instrument.py @@ -56,6 +56,46 @@ class StandaloneCallMatch: ) +def is_inside_string(code: str, pos: int) -> bool: + """Check if a position in code is inside a string literal. + + Handles single quotes, double quotes, and template literals (backticks). + Properly handles escaped quotes. + + Args: + code: The source code. + pos: The position to check. + + Returns: + True if the position is inside a string literal. + + """ + in_string = False + string_char = None + i = 0 + + while i < pos: + char = code[i] + + if in_string: + # Check for escape sequence + if char == "\\" and i + 1 < len(code): + i += 2 # Skip escaped character + continue + # Check for end of string + if char == string_char: + in_string = False + string_char = None + # Check for start of string + elif char in "\"'`": + in_string = True + string_char = char + + i += 1 + + return in_string + + class StandaloneCallTransformer: """Transforms standalone func(...) calls in JavaScript test code. @@ -82,6 +122,11 @@ def __init__(self, function_to_optimize: FunctionToOptimize, capture_func: str) # Captures: (whitespace)(await )?(object.)*func_name( # We'll filter out expect() and codeflash. cases in the transform loop self._call_pattern = re.compile(rf"(\s*)(await\s+)?((?:\w+\.)*){re.escape(self.func_name)}\s*\(") + # Pattern to match bracket notation: obj['func_name']( or obj["func_name"]( + # Captures: (whitespace)(await )?(obj)['|"]func_name['|"]( + self._bracket_call_pattern = re.compile( + rf"(\s*)(await\s+)?(\w+)\[['\"]({re.escape(self.func_name)})['\"]]\s*\(" + ) def transform(self, code: str) -> str: """Transform all standalone calls in the code.""" @@ -89,7 +134,25 @@ def transform(self, code: str) -> str: pos = 0 while pos < len(code): - match = self._call_pattern.search(code, pos) + # Try both dot notation and bracket notation patterns + dot_match = self._call_pattern.search(code, pos) + bracket_match = self._bracket_call_pattern.search(code, pos) + + # Choose the first match (by position) + match = None + is_bracket_notation = False + if dot_match and bracket_match: + if dot_match.start() <= bracket_match.start(): + match = dot_match + else: + match = bracket_match + is_bracket_notation = True + elif dot_match: + match = dot_match + elif bracket_match: + match = bracket_match + is_bracket_notation = True + if not match: result.append(code[pos:]) break @@ -106,7 +169,11 @@ def transform(self, code: str) -> str: result.append(code[pos:match_start]) # Try to parse the full standalone call - standalone_match = self._parse_standalone_call(code, match) + if is_bracket_notation: + standalone_match = self._parse_bracket_standalone_call(code, match) + else: + standalone_match = self._parse_standalone_call(code, match) + if standalone_match is None: # Couldn't parse, skip this match result.append(code[match_start : match.end()]) @@ -115,7 +182,7 @@ def transform(self, code: str) -> str: # Generate the transformed code self.invocation_counter += 1 - transformed = self._generate_transformed_call(standalone_match) + transformed = self._generate_transformed_call(standalone_match, is_bracket_notation) result.append(transformed) pos = standalone_match.end_pos @@ -123,6 +190,10 @@ def transform(self, code: str) -> str: def _should_skip_match(self, code: str, start: int, match: re.Match) -> bool: """Check if the match should be skipped (inside expect, already transformed, etc.).""" + # Skip if inside a string literal (e.g., test description) + if is_inside_string(code, start): + return True + # Look backwards to check context lookback_start = max(0, start - 200) lookback = code[lookback_start:start] @@ -252,17 +323,24 @@ def _find_balanced_parens(self, code: str, open_paren_pos: int) -> tuple[str | N in_string = False string_char = None - while pos < len(code) and depth > 0: - char = code[pos] + s = code # local alias for speed + s_len = len(s) + quotes = "\"'`" + + while pos < s_len and depth > 0: + char = s[pos] # Handle string literals - if char in "\"'`" and (pos == 0 or code[pos - 1] != "\\"): - if not in_string: - in_string = True - string_char = char - elif char == string_char: - in_string = False - string_char = None + # Note: preserve original escaping semantics (only checks immediate preceding char) + if char in quotes: + prev_char = s[pos - 1] if pos > 0 else None + if prev_char != "\\": + if not in_string: + in_string = True + string_char = char + elif char == string_char: + in_string = False + string_char = None elif not in_string: if char == "(": depth += 1 @@ -274,19 +352,64 @@ def _find_balanced_parens(self, code: str, open_paren_pos: int) -> tuple[str | N if depth != 0: return None, -1 - return code[open_paren_pos + 1 : pos - 1], pos + # slice once + return s[open_paren_pos + 1 : pos - 1], pos - def _generate_transformed_call(self, match: StandaloneCallMatch) -> str: + def _parse_bracket_standalone_call(self, code: str, match: re.Match) -> StandaloneCallMatch | None: + """Parse a complete standalone obj['func'](...) call with bracket notation.""" + leading_ws = match.group(1) + prefix = match.group(2) or "" # "await " or "" + obj_name = match.group(3) # The object name before bracket + # match.group(4) is the function name inside brackets + + # Find the opening paren position + match_text = match.group(0) + paren_offset = match_text.rfind("(") + open_paren_pos = match.start() + paren_offset + + # Find the arguments (content inside parens) + func_args, close_pos = self._find_balanced_parens(code, open_paren_pos) + if func_args is None: + return None + + # Check for trailing semicolon + end_pos = close_pos + # Skip whitespace + s = code + s_len = len(s) + while end_pos < s_len and s[end_pos] in " \t": + end_pos += 1 + + has_trailing_semicolon = end_pos < s_len and s[end_pos] == ";" + if has_trailing_semicolon: + end_pos += 1 + + return StandaloneCallMatch( + start_pos=match.start(), + end_pos=end_pos, + leading_whitespace=leading_ws, + func_args=func_args, + prefix=prefix, + object_prefix=f"{obj_name}.", # Use dot notation format for consistency + has_trailing_semicolon=has_trailing_semicolon, + ) + + def _generate_transformed_call(self, match: StandaloneCallMatch, is_bracket_notation: bool = False) -> str: """Generate the transformed code for a standalone call.""" line_id = str(self.invocation_counter) args_str = match.func_args.strip() semicolon = ";" if match.has_trailing_semicolon else "" - # Handle method calls on objects (e.g., calc.fibonacci, this.method) + # Handle method calls on objects (e.g., calc.fibonacci, this.method, instance['method']) if match.object_prefix: # Remove trailing dot from object prefix for the bind call obj = match.object_prefix.rstrip(".") - full_method = f"{obj}.{self.func_name}" + + # For bracket notation, use bracket access syntax for the bind + if is_bracket_notation: + full_method = f"{obj}['{self.func_name}']" + else: + full_method = f"{obj}.{self.func_name}" if args_str: return ( @@ -370,6 +493,12 @@ def transform(self, code: str) -> str: result.append(code[pos:]) break + # Skip if inside a string literal (e.g., test description) + if is_inside_string(code, match.start()): + result.append(code[pos : match.end()]) + pos = match.end() + continue + # Add everything before the match result.append(code[pos : match.start()]) @@ -792,7 +921,7 @@ def validate_and_fix_import_style(test_code: str, source_file_path: Path, functi Fixed test code with correct import style. """ - from codeflash.languages.treesitter_utils import get_analyzer_for_file + from codeflash.languages.javascript.treesitter import get_analyzer_for_file # Read source file to determine export style try: @@ -901,6 +1030,115 @@ def is_relevant_import(module_path: str) -> bool: return test_code +def fix_import_path_for_test_location( + test_code: str, source_file_path: Path, test_file_path: Path, module_root: Path +) -> str: + """Fix import paths in generated test code to be relative to test file location. + + The AI may generate tests with import paths that are relative to the module root + (e.g., 'apps/web/app/file') instead of relative to where the test file is located + (e.g., '../../app/file'). This function fixes such imports. + + Args: + test_code: The generated test code. + source_file_path: Absolute path to the source file being tested. + test_file_path: Absolute path to where the test file will be written. + module_root: Root directory of the module/project. + + Returns: + Test code with corrected import paths. + + """ + import os + + # Calculate the correct relative import path from test file to source file + test_dir = test_file_path.parent + try: + correct_rel_path = os.path.relpath(source_file_path, test_dir) + correct_rel_path = correct_rel_path.replace("\\", "/") + # Remove file extension for JS/TS imports + for ext in [".tsx", ".ts", ".jsx", ".js", ".mjs", ".cjs"]: + if correct_rel_path.endswith(ext): + correct_rel_path = correct_rel_path[: -len(ext)] + break + # Ensure it starts with ./ or ../ + if not correct_rel_path.startswith("."): + correct_rel_path = "./" + correct_rel_path + except ValueError: + # Can't compute relative path (different drives on Windows) + return test_code + + # Try to compute what incorrect path the AI might have generated + # The AI often uses module_root-relative paths like 'apps/web/app/...' + try: + source_rel_to_module = os.path.relpath(source_file_path, module_root) + source_rel_to_module = source_rel_to_module.replace("\\", "/") + # Remove extension + for ext in [".tsx", ".ts", ".jsx", ".js", ".mjs", ".cjs"]: + if source_rel_to_module.endswith(ext): + source_rel_to_module = source_rel_to_module[: -len(ext)] + break + except ValueError: + return test_code + + # Also check for project root-relative paths (including module_root in path) + try: + project_root = module_root.parent if module_root.name in ["src", "lib", "app", "web", "apps"] else module_root + source_rel_to_project = os.path.relpath(source_file_path, project_root) + source_rel_to_project = source_rel_to_project.replace("\\", "/") + for ext in [".tsx", ".ts", ".jsx", ".js", ".mjs", ".cjs"]: + if source_rel_to_project.endswith(ext): + source_rel_to_project = source_rel_to_project[: -len(ext)] + break + except ValueError: + source_rel_to_project = None + + # Source file name (for matching module paths that end with the file name) + source_name = source_file_path.stem + + # Patterns to find import statements + # ESM: import { func } from 'path' or import func from 'path' + esm_import_pattern = re.compile(r"(import\s+(?:{[^}]+}|\w+)\s+from\s+['\"])([^'\"]+)(['\"])") + # CommonJS: const { func } = require('path') or const func = require('path') + cjs_require_pattern = re.compile( + r"((?:const|let|var)\s+(?:{[^}]+}|\w+)\s*=\s*require\s*\(\s*['\"])([^'\"]+)(['\"])" + ) + + def should_fix_path(import_path: str) -> bool: + """Check if this import path looks like it should point to our source file.""" + # Skip relative imports that already look correct + if import_path.startswith(("./", "../")): + return False + # Skip package imports (no path separators or start with @) + if "/" not in import_path and "\\" not in import_path: + return False + if import_path.startswith("@") and "/" in import_path: + # Could be an alias like @/utils - skip these + return False + # Check if it looks like it points to our source file + if import_path == source_rel_to_module: + return True + if source_rel_to_project and import_path == source_rel_to_project: + return True + if import_path.endswith((source_name, "/" + source_name)): + return True + return False + + def fix_import(match: re.Match[str]) -> str: + """Replace incorrect import path with correct relative path.""" + prefix = match.group(1) + import_path = match.group(2) + suffix = match.group(3) + + if should_fix_path(import_path): + logger.debug(f"Fixing import path: {import_path} -> {correct_rel_path}") + return f"{prefix}{correct_rel_path}{suffix}" + return match.group(0) + + test_code = esm_import_pattern.sub(fix_import, test_code) + return cjs_require_pattern.sub(fix_import, test_code) + + def get_instrumented_test_path(original_path: Path, mode: str) -> Path: """Generate path for instrumented test file. @@ -962,3 +1200,171 @@ def instrument_generated_js_test( mode=mode, remove_assertions=True, ) + + +def fix_imports_inside_test_blocks(test_code: str) -> str: + """Fix import statements that appear inside test/it blocks. + + JavaScript/TypeScript `import` statements must be at the top level of a module. + The AI sometimes generates imports inside test functions, which is invalid syntax. + + This function detects such patterns and converts them to dynamic require() calls + which are valid inside functions. + + Args: + test_code: The generated test code. + + Returns: + Fixed test code with imports converted to require() inside functions. + + """ + if not test_code or not test_code.strip(): + return test_code + + # Pattern to match import statements inside functions + # This captures imports that appear after function/test block openings + # We look for lines that: + # 1. Start with whitespace (indicating they're inside a block) + # 2. Have an import statement + + lines = test_code.split("\n") + result_lines = [] + brace_depth = 0 + in_test_block = False + + for line in lines: + stripped = line.strip() + + # Track brace depth to know if we're inside a block + # Count braces, but ignore braces in strings (simplified check) + for char in stripped: + if char == "{": + brace_depth += 1 + elif char == "}": + brace_depth -= 1 + + # Check if we're entering a test/it/describe block + if re.match(r"^(test|it|describe|beforeEach|afterEach|beforeAll|afterAll)\s*\(", stripped): + in_test_block = True + + # Check for import statement inside a block (brace_depth > 0 means we're inside a function/block) + if brace_depth > 0 and stripped.startswith("import "): + # Convert ESM import to require + # Pattern: import { name } from 'module' -> const { name } = require('module') + # Pattern: import name from 'module' -> const name = require('module') + + named_import = re.match(r"import\s+\{([^}]+)\}\s+from\s+['\"]([^'\"]+)['\"]", stripped) + default_import = re.match(r"import\s+(\w+)\s+from\s+['\"]([^'\"]+)['\"]", stripped) + namespace_import = re.match(r"import\s+\*\s+as\s+(\w+)\s+from\s+['\"]([^'\"]+)['\"]", stripped) + + leading_whitespace = line[: len(line) - len(line.lstrip())] + + if named_import: + names = named_import.group(1) + module = named_import.group(2) + new_line = f"{leading_whitespace}const {{{names}}} = require('{module}');" + result_lines.append(new_line) + logger.debug(f"Fixed import inside block: {stripped} -> {new_line.strip()}") + continue + if default_import: + name = default_import.group(1) + module = default_import.group(2) + new_line = f"{leading_whitespace}const {name} = require('{module}');" + result_lines.append(new_line) + logger.debug(f"Fixed import inside block: {stripped} -> {new_line.strip()}") + continue + if namespace_import: + name = namespace_import.group(1) + module = namespace_import.group(2) + new_line = f"{leading_whitespace}const {name} = require('{module}');" + result_lines.append(new_line) + logger.debug(f"Fixed import inside block: {stripped} -> {new_line.strip()}") + continue + + result_lines.append(line) + + return "\n".join(result_lines) + + +def fix_jest_mock_paths(test_code: str, test_file_path: Path, source_file_path: Path, tests_root: Path) -> str: + """Fix relative paths in jest.mock() calls to be correct from the test file's location. + + The AI sometimes generates jest.mock() calls with paths relative to the source file + instead of the test file. For example: + - Source at `src/queue/queue.ts` imports `../environment` (-> src/environment) + - Test at `tests/test.test.ts` generates `jest.mock('../environment')` (-> ./environment, wrong!) + - Should generate `jest.mock('../src/environment')` + + This function detects relative mock paths and adjusts them based on the test file's + location relative to the source file's directory. + + Args: + test_code: The generated test code. + test_file_path: Path to the test file being generated. + source_file_path: Path to the source file being tested. + tests_root: Root directory of the tests. + + Returns: + Fixed test code with corrected mock paths. + + """ + if not test_code or not test_code.strip(): + return test_code + + import os + + # Get the directory containing the source file and the test file + source_dir = source_file_path.resolve().parent + test_dir = test_file_path.resolve().parent + project_root = tests_root.resolve().parent if tests_root.name == "tests" else tests_root.resolve() + + # Pattern to match jest.mock() or jest.doMock() with relative paths + mock_pattern = re.compile(r"(jest\.(?:mock|doMock)\s*\(\s*['\"])(\.\./[^'\"]+|\.\/[^'\"]+)(['\"])") + + def fix_mock_path(match: re.Match[str]) -> str: + original = match.group(0) + prefix = match.group(1) + rel_path = match.group(2) + suffix = match.group(3) + + # Resolve the path as if it were relative to the source file's directory + # (which is how the AI often generates it) + source_relative_resolved = (source_dir / rel_path).resolve() + + # Check if this resolved path exists or if adjusting it would make more sense + # Calculate what the correct relative path from the test file should be + try: + # First, try to find if the path makes sense from the test directory + test_relative_resolved = (test_dir / rel_path).resolve() + + # If the path exists relative to test dir, keep it + if test_relative_resolved.exists() or ( + test_relative_resolved.with_suffix(".ts").exists() + or test_relative_resolved.with_suffix(".js").exists() + or test_relative_resolved.with_suffix(".tsx").exists() + or test_relative_resolved.with_suffix(".jsx").exists() + ): + return original # Keep original, it's valid + + # If path exists relative to source dir, recalculate from test dir + if source_relative_resolved.exists() or ( + source_relative_resolved.with_suffix(".ts").exists() + or source_relative_resolved.with_suffix(".js").exists() + or source_relative_resolved.with_suffix(".tsx").exists() + or source_relative_resolved.with_suffix(".jsx").exists() + ): + # Calculate the correct relative path from test_dir to source_relative_resolved + new_rel_path = Path(os.path.relpath(source_relative_resolved, test_dir)).as_posix() + # Ensure it starts with ./ or ../ + if not new_rel_path.startswith("../") and not new_rel_path.startswith("./"): + new_rel_path = f"./{new_rel_path}" + + logger.debug(f"Fixed jest.mock path: {rel_path} -> {new_rel_path}") + return f"{prefix}{new_rel_path}{suffix}" + + except (ValueError, OSError): + pass # Path resolution failed, keep original + + return original # Keep original if we can't fix it + + return mock_pattern.sub(fix_mock_path, test_code) diff --git a/codeflash/languages/javascript/line_profiler.py b/codeflash/languages/javascript/line_profiler.py index 57f046d4a..81b38983c 100644 --- a/codeflash/languages/javascript/line_profiler.py +++ b/codeflash/languages/javascript/line_profiler.py @@ -11,7 +11,7 @@ import logging from typing import TYPE_CHECKING -from codeflash.languages.treesitter_utils import get_analyzer_for_file +from codeflash.languages.javascript.treesitter import get_analyzer_for_file if TYPE_CHECKING: from pathlib import Path diff --git a/codeflash/languages/javascript/module_system.py b/codeflash/languages/javascript/module_system.py index 3e3ff29dc..89d723c02 100644 --- a/codeflash/languages/javascript/module_system.py +++ b/codeflash/languages/javascript/module_system.py @@ -100,23 +100,40 @@ def detect_module_system(project_root: Path, file_path: Path | None = None) -> s try: content = file_path.read_text() - # Look for ES module syntax + # Look for ES module syntax - these are explicit ESM markers has_import = "import " in content and "from " in content - has_export = "export " in content or "export default" in content or "export {" in content + # Check for export function/class/const/default which are unambiguous ESM syntax + has_esm_export = ( + "export function " in content + or "export class " in content + or "export const " in content + or "export let " in content + or "export default " in content + or "export async function " in content + ) + has_export_block = "export {" in content # Look for CommonJS syntax has_require = "require(" in content has_module_exports = "module.exports" in content or "exports." in content - # Determine based on what we found - if (has_import or has_export) and not (has_require or has_module_exports): - logger.debug("Detected ES Module from import/export statements") + # Prioritize ESM when explicit ESM export syntax is found + # This handles hybrid files that have both `export function` and `module.exports` + # The ESM syntax is more explicit and should take precedence + if has_esm_export or has_import: + logger.debug("Detected ES Module from explicit export/import statements") return ModuleSystem.ES_MODULE - if (has_require or has_module_exports) and not (has_import or has_export): + # Pure CommonJS + if (has_require or has_module_exports) and not has_export_block: logger.debug("Detected CommonJS from require/module.exports") return ModuleSystem.COMMONJS + # Export block without other ESM markers - still ESM + if has_export_block: + logger.debug("Detected ES Module from export block") + return ModuleSystem.ES_MODULE + except Exception as e: logger.warning("Failed to analyze file %s: %s", file_path, e) @@ -416,8 +433,10 @@ def ensure_module_system_compatibility(code: str, target_module_system: str, pro is_esm = has_import or has_export # Convert if needed - if target_module_system == ModuleSystem.ES_MODULE and is_commonjs and not is_esm: - logger.debug("Converting CommonJS to ES Module syntax") + # For ESM target: convert any require statements, even if there are also import statements + # This handles generated tests that have ESM imports for test globals but CommonJS for the function + if target_module_system == ModuleSystem.ES_MODULE and has_require: + logger.debug("Converting CommonJS require statements to ES Module syntax") return convert_commonjs_to_esm(code) if target_module_system == ModuleSystem.COMMONJS and is_esm and not is_commonjs: diff --git a/codeflash/languages/javascript/parse.py b/codeflash/languages/javascript/parse.py index d6b43feae..e3eee4831 100644 --- a/codeflash/languages/javascript/parse.py +++ b/codeflash/languages/javascript/parse.py @@ -122,6 +122,7 @@ def parse_jest_test_xml( # This handles cases where instrumented files are in temp directories instrumented_path_lookup: dict[str, tuple[Path, TestType]] = {} for test_file in test_files.test_files: + # Add behavior instrumented file paths if test_file.instrumented_behavior_file_path: # Store both the absolute path and resolved path as keys abs_path = str(test_file.instrumented_behavior_file_path.resolve()) @@ -132,18 +133,35 @@ def parse_jest_test_xml( test_file.test_type, ) logger.debug(f"Jest XML lookup: registered {abs_path}") + # Also add benchmarking file paths (perf-only instrumented tests) + if test_file.benchmarking_file_path: + bench_abs_path = str(test_file.benchmarking_file_path.resolve()) + instrumented_path_lookup[bench_abs_path] = (test_file.benchmarking_file_path, test_file.test_type) + instrumented_path_lookup[str(test_file.benchmarking_file_path)] = ( + test_file.benchmarking_file_path, + test_file.test_type, + ) + logger.debug(f"Jest XML lookup: registered benchmark {bench_abs_path}") # Also build a filename-only lookup for fallback matching # This handles cases where JUnit XML has relative paths that don't match absolute paths # e.g., JUnit has "test/utils__perfinstrumented.test.ts" but lookup has absolute paths filename_lookup: dict[str, tuple[Path, TestType]] = {} for test_file in test_files.test_files: + # Add instrumented_behavior_file_path (behavior tests) if test_file.instrumented_behavior_file_path: filename = test_file.instrumented_behavior_file_path.name # Only add if not already present (avoid overwrites in case of duplicate filenames) if filename not in filename_lookup: filename_lookup[filename] = (test_file.instrumented_behavior_file_path, test_file.test_type) logger.debug(f"Jest XML filename lookup: registered {filename}") + # Also add benchmarking_file_path (perf-only tests) - these have different filenames + # e.g., utils__perfonlyinstrumented.test.ts vs utils__perfinstrumented.test.ts + if test_file.benchmarking_file_path: + bench_filename = test_file.benchmarking_file_path.name + if bench_filename not in filename_lookup: + filename_lookup[bench_filename] = (test_file.benchmarking_file_path, test_file.test_type) + logger.debug(f"Jest XML filename lookup: registered benchmark file {bench_filename}") # Fallback: if JUnit XML doesn't have system-out, use subprocess stdout directly global_stdout = "" @@ -157,6 +175,19 @@ def parse_jest_test_xml( logger.debug(f"Found {marker_count} timing start markers in Jest stdout") else: logger.debug(f"No timing start markers found in Jest stdout (len={len(global_stdout)})") + # Check for END markers with duration (perf test markers) + end_marker_count = len(jest_end_pattern.findall(global_stdout)) + if end_marker_count > 0: + logger.debug( + f"[PERF-DEBUG] Found {end_marker_count} END timing markers with duration in Jest stdout" + ) + # Sample a few markers to verify loop indices + end_samples = list(jest_end_pattern.finditer(global_stdout))[:5] + for sample in end_samples: + groups = sample.groups() + logger.debug(f"[PERF-DEBUG] Sample END marker: loopIndex={groups[3]}, duration={groups[5]}") + else: + logger.debug("[PERF-DEBUG] No END markers with duration found in Jest stdout") except (AttributeError, UnicodeDecodeError): global_stdout = "" @@ -167,18 +198,47 @@ def parse_jest_test_xml( # Extract console output from suite-level system-out (Jest specific) suite_stdout = _extract_jest_console_output(suite._elem) # noqa: SLF001 - # Fallback: use subprocess stdout if XML system-out is empty - if not suite_stdout and global_stdout: - suite_stdout = global_stdout + # Combine suite stdout with global stdout to ensure we capture all timing markers + # Jest-junit may not capture all console.log output in the XML, so we also need + # to check the subprocess stdout directly for timing markers + combined_stdout = suite_stdout + if global_stdout: + if combined_stdout: + combined_stdout = combined_stdout + "\n" + global_stdout + else: + combined_stdout = global_stdout - # Parse timing markers from the suite's console output - start_matches = list(jest_start_pattern.finditer(suite_stdout)) + # Parse timing markers from the combined console output + start_matches = list(jest_start_pattern.finditer(combined_stdout)) end_matches_dict = {} - for match in jest_end_pattern.finditer(suite_stdout): + for match in jest_end_pattern.finditer(combined_stdout): # Key: (testName, testName2, funcName, loopIndex, lineId) key = match.groups()[:5] end_matches_dict[key] = match + # Debug: log suite-level END marker parsing for perf tests + if end_matches_dict: + # Get unique loop indices from the parsed END markers + loop_indices = sorted({int(k[3]) if k[3].isdigit() else 1 for k in end_matches_dict}) + logger.debug( + f"[PERF-DEBUG] Suite {suite_count}: parsed {len(end_matches_dict)} END markers from suite_stdout, loop_index range: {min(loop_indices)}-{max(loop_indices)}" + ) + + # Also collect timing markers from testcase-level system-out (Vitest puts output at testcase level) + for tc in suite: + tc_system_out = tc._elem.find("system-out") # noqa: SLF001 + if tc_system_out is not None and tc_system_out.text: + tc_stdout = tc_system_out.text.strip() + logger.debug(f"Vitest testcase system-out found: {len(tc_stdout)} chars, first 200: {tc_stdout[:200]}") + end_marker_count = 0 + for match in jest_end_pattern.finditer(tc_stdout): + key = match.groups()[:5] + end_matches_dict[key] = match + end_marker_count += 1 + if end_marker_count > 0: + logger.debug(f"Found {end_marker_count} END timing markers in testcase system-out") + start_matches.extend(jest_start_pattern.finditer(tc_stdout)) + for testcase in suite: testcase_count += 1 test_class_path = testcase.classname # For Jest, this is the file path @@ -264,7 +324,7 @@ def parse_jest_test_xml( # Infer test type from filename pattern filename = test_file_path.name if "__perf_test_" in filename or "_perf_test_" in filename: - test_type = TestType.GENERATED_PERFORMANCE + test_type = TestType.GENERATED_REGRESSION # Performance tests are still generated regression tests elif "__unit_test_" in filename or "_unit_test_" in filename: test_type = TestType.GENERATED_REGRESSION else: @@ -294,6 +354,13 @@ def parse_jest_test_xml( sanitized_test_name = re.sub(r"[!#: ()\[\]{}|\\/*?^$.+\-]", "_", test_name) matching_starts = [m for m in start_matches if sanitized_test_name in m.group(2)] + # Debug: log which branch we're taking + logger.debug( + f"[FLOW-DEBUG] Testcase '{test_name[:50]}': " + f"total_start_matches={len(start_matches)}, matching_starts={len(matching_starts)}, " + f"total_end_matches={len(end_matches_dict)}" + ) + # For performance tests (capturePerf), there are no START markers - only END markers with duration # Check for END markers directly if no START markers found matching_ends_direct = [] @@ -304,9 +371,42 @@ def parse_jest_test_xml( # end_key is (module, testName, funcName, loopIndex, invocationId) if len(end_key) >= 2 and sanitized_test_name in end_key[1]: matching_ends_direct.append(end_match) + # Debug: log matching results for perf tests + if matching_ends_direct: + loop_indices = [int(m.groups()[3]) if m.groups()[3].isdigit() else 1 for m in matching_ends_direct] + logger.debug( + f"[PERF-MATCH] Testcase '{test_name[:40]}': matched {len(matching_ends_direct)} END markers, " + f"loop_index range: {min(loop_indices)}-{max(loop_indices)}" + ) + elif end_matches_dict: + # No matches but we have END markers - check why + sample_keys = list(end_matches_dict.keys())[:3] + logger.debug( + f"[PERF-MISMATCH] Testcase '{test_name[:40]}': no matches found. " + f"sanitized_test_name='{sanitized_test_name[:50]}', " + f"sample end_keys={[k[1][:30] if len(k) >= 2 else k for k in sample_keys]}" + ) + + # Log if we're skipping the matching_ends_direct branch + if matching_starts and end_matches_dict: + logger.debug( + f"[FLOW-SKIP] Testcase '{test_name[:40]}': has {len(matching_starts)} START markers, " + f"skipping {len(end_matches_dict)} END markers (behavior test mode)" + ) if not matching_starts and not matching_ends_direct: - # No timing markers found - add basic result + # No timing markers found - use JUnit XML time attribute as fallback + # The time attribute is in seconds (e.g., "0.00077875"), convert to nanoseconds + runtime = None + try: + time_attr = testcase._elem.attrib.get("time") # noqa: SLF001 + if time_attr: + time_seconds = float(time_attr) + runtime = int(time_seconds * 1_000_000_000) # Convert seconds to nanoseconds + logger.debug(f"Jest XML: using time attribute for {test_name}: {time_seconds}s = {runtime}ns") + except (ValueError, TypeError) as e: + logger.debug(f"Jest XML: could not parse time attribute: {e}") + test_results.add( FunctionTestInvocation( loop_index=1, @@ -318,7 +418,7 @@ def parse_jest_test_xml( iteration_id="", ), file_name=test_file_path, - runtime=None, + runtime=runtime, test_framework=test_config.test_framework, did_pass=result, test_type=test_type, @@ -329,11 +429,13 @@ def parse_jest_test_xml( ) elif matching_ends_direct: # Performance test format: process END markers directly (no START markers) + loop_indices_found = [] for end_match in matching_ends_direct: groups = end_match.groups() # groups: (module, testName, funcName, loopIndex, invocationId, durationNs) func_name = groups[2] loop_index = int(groups[3]) if groups[3].isdigit() else 1 + loop_indices_found.append(loop_index) line_id = groups[4] try: runtime = int(groups[5]) @@ -359,6 +461,12 @@ def parse_jest_test_xml( stdout="", ) ) + if loop_indices_found: + logger.debug( + f"[LOOP-DEBUG] Testcase '{test_name}': processed {len(matching_ends_direct)} END markers, " + f"loop_index range: {min(loop_indices_found)}-{max(loop_indices_found)}, " + f"total results so far: {len(test_results.test_results)}" + ) else: # Process each timing marker for match in matching_starts: @@ -410,5 +518,14 @@ def parse_jest_test_xml( f"Jest XML parsing complete: {len(test_results.test_results)} results " f"from {suite_count} suites, {testcase_count} testcases" ) + # Debug: show loop_index distribution for perf analysis + if test_results.test_results: + loop_indices = [r.loop_index for r in test_results.test_results] + unique_loop_indices = sorted(set(loop_indices)) + min_idx, max_idx = min(unique_loop_indices), max(unique_loop_indices) + logger.debug( + f"[LOOP-SUMMARY] Results loop_index: min={min_idx}, max={max_idx}, " + f"unique_count={len(unique_loop_indices)}, total_results={len(loop_indices)}" + ) return test_results diff --git a/codeflash/languages/javascript/support.py b/codeflash/languages/javascript/support.py index 7ba69ce50..e0111c634 100644 --- a/codeflash/languages/javascript/support.py +++ b/codeflash/languages/javascript/support.py @@ -14,15 +14,16 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import CodeContext, FunctionFilterCriteria, HelperFunction, Language, TestInfo, TestResult +from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage, get_analyzer_for_file from codeflash.languages.registry import register_language -from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage, get_analyzer_for_file from codeflash.models.models import FunctionParent if TYPE_CHECKING: from collections.abc import Sequence from codeflash.languages.base import ReferenceInfo - from codeflash.languages.treesitter_utils import TypeDefinition + from codeflash.languages.javascript.treesitter import TypeDefinition + from codeflash.models.models import GeneratedTestsList, InvocationId logger = logging.getLogger(__name__) @@ -63,6 +64,10 @@ def test_framework(self) -> str: def comment_prefix(self) -> str: return "//" + @property + def dir_excludes(self) -> frozenset[str]: + return frozenset({"node_modules", "dist", "build", ".next", ".nuxt", "coverage", ".cache", ".turbo", ".vercel"}) + # === Discovery === def discover_functions( @@ -104,6 +109,12 @@ def discover_functions( if not criteria.include_async and func.is_async: continue + # Skip non-exported functions (can't be imported in tests) + # Exception: nested functions and methods are allowed if their parent is exported + if not func.is_exported and not func.parent_function: + logger.debug(f"Skipping non-exported function: {func.name}") # noqa: G004 + continue + # Build parents list parents: list[FunctionParent] = [] if func.class_name: @@ -326,8 +337,14 @@ def extract_code_context(self, function: FunctionToOptimize, project_root: Path, else: target_code = "" + imports = analyzer.find_imports(source) + + # Find helper functions called by target (needed before class wrapping to find same-class helpers) + helpers = self._find_helper_functions(function, source, analyzer, imports, module_root) + # For class methods, wrap the method in its class definition # This is necessary because method definition syntax is only valid inside a class body + same_class_helper_names: set[str] = set() if function.is_method and function.parents: class_name = None for parent in function.parents: @@ -336,17 +353,26 @@ def extract_code_context(self, function: FunctionToOptimize, project_root: Path, break if class_name: + # Find same-class helper methods that need to be included inside the class wrapper + same_class_helpers = self._find_same_class_helpers( + class_name, function.function_name, helpers, tree_functions, lines + ) + same_class_helper_names = {h[0] for h in same_class_helpers} # method names + # Find the class definition in the source to get proper indentation, JSDoc, constructor, and fields class_info = self._find_class_definition(source, class_name, analyzer, function.function_name) if class_info: class_jsdoc, class_indent, constructor_code, fields_code = class_info - # Build the class body with fields, constructor, and target method + # Build the class body with fields, constructor, target method, and same-class helpers class_body_parts = [] if fields_code: class_body_parts.append(fields_code) if constructor_code: class_body_parts.append(constructor_code) class_body_parts.append(target_code) + # Add same-class helper methods inside the class body + for _helper_name, helper_source in same_class_helpers: + class_body_parts.append(helper_source) class_body = "\n".join(class_body_parts) # Wrap the method in a class definition with context @@ -357,13 +383,16 @@ def extract_code_context(self, function: FunctionToOptimize, project_root: Path, else: target_code = f"{class_indent}class {class_name} {{\n{class_body}{class_indent}}}\n" else: - # Fallback: wrap with no indentation - target_code = f"class {class_name} {{\n{target_code}}}\n" - - imports = analyzer.find_imports(source) + # Fallback: wrap with no indentation, including same-class helpers + helper_code = "\n".join(h[1] for h in same_class_helpers) + if helper_code: + target_code = f"class {class_name} {{\n{target_code}\n{helper_code}}}\n" + else: + target_code = f"class {class_name} {{\n{target_code}}}\n" - # Find helper functions called by target - helpers = self._find_helper_functions(function, source, analyzer, imports, module_root) + # Filter out same-class helpers from the helpers list (they're already inside the class wrapper) + if same_class_helper_names: + helpers = [h for h in helpers if h.name not in same_class_helper_names] # Extract import statements as strings import_lines = [] @@ -546,6 +575,49 @@ def _extract_class_context( return (constructor_code, fields_code) + def _find_same_class_helpers( + self, + class_name: str, + target_method_name: str, + helpers: list[HelperFunction], + tree_functions: list, + lines: list[str], + ) -> list[tuple[str, str]]: + """Find helper methods that belong to the same class as the target method. + + These helpers need to be included inside the class wrapper rather than + appended outside, because they may use class-specific syntax like 'private'. + + Args: + class_name: Name of the class containing the target method. + target_method_name: Name of the target method (to exclude). + helpers: List of all helper functions found. + tree_functions: List of FunctionNode from tree-sitter analysis. + lines: Source code split into lines. + + Returns: + List of (method_name, source_code) tuples for same-class helpers. + + """ + same_class_helpers: list[tuple[str, str]] = [] + + # Build a set of helper names for quick lookup + helper_names = {h.name for h in helpers} + + # Names to exclude from same-class helpers (target method and constructor) + exclude_names = {target_method_name, "constructor"} + + # Find methods in tree_functions that belong to the same class and are helpers + for func in tree_functions: + if func.class_name == class_name and func.name in helper_names and func.name not in exclude_names: + # Extract source including JSDoc if present + effective_start = func.doc_start_line or func.start_line + helper_lines = lines[effective_start - 1 : func.end_line] + helper_source = "".join(helper_lines) + same_class_helpers.append((func.name, helper_source)) + + return same_class_helpers + def _find_helper_functions( self, function: FunctionToOptimize, @@ -1707,6 +1779,116 @@ def remove_test_functions(self, test_source: str, functions_to_remove: list[str] return remove_test_functions(test_source, functions_to_remove) + def postprocess_generated_tests( + self, generated_tests: GeneratedTestsList, test_framework: str, project_root: Path, source_file_path: Path + ) -> GeneratedTestsList: + """Apply language-specific postprocessing to generated tests.""" + from codeflash.languages.javascript.edit_tests import ( + disable_ts_check, + inject_test_globals, + normalize_generated_tests_imports, + ) + from codeflash.languages.javascript.module_system import detect_module_system + + module_system = detect_module_system(project_root, source_file_path) + if module_system == "esm": + generated_tests = inject_test_globals(generated_tests, test_framework) + if self.language == Language.TYPESCRIPT: + generated_tests = disable_ts_check(generated_tests) + return normalize_generated_tests_imports(generated_tests) + + def remove_test_functions_from_generated_tests( + self, generated_tests: GeneratedTestsList, functions_to_remove: list[str] + ) -> GeneratedTestsList: + """Remove specific test functions from generated tests.""" + from codeflash.models.models import GeneratedTests, GeneratedTestsList + + updated_tests: list[GeneratedTests] = [] + for test in generated_tests.generated_tests: + updated_tests.append( + GeneratedTests( + generated_original_test_source=self.remove_test_functions( + test.generated_original_test_source, functions_to_remove + ), + instrumented_behavior_test_source=test.instrumented_behavior_test_source, + instrumented_perf_test_source=test.instrumented_perf_test_source, + behavior_file_path=test.behavior_file_path, + perf_file_path=test.perf_file_path, + ) + ) + return GeneratedTestsList(generated_tests=updated_tests) + + def add_runtime_comments_to_generated_tests( + self, + generated_tests: GeneratedTestsList, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], + tests_project_rootdir: Path | None = None, + ) -> GeneratedTestsList: + """Add runtime comments to generated tests.""" + from codeflash.models.models import GeneratedTests, GeneratedTestsList + + tests_root = tests_project_rootdir or Path() + original_runtimes_dict = self._build_runtime_map(original_runtimes, tests_root) + optimized_runtimes_dict = self._build_runtime_map(optimized_runtimes, tests_root) + + modified_tests: list[GeneratedTests] = [] + for test in generated_tests.generated_tests: + modified_source = self.add_runtime_comments( + test.generated_original_test_source, original_runtimes_dict, optimized_runtimes_dict + ) + modified_tests.append( + GeneratedTests( + generated_original_test_source=modified_source, + instrumented_behavior_test_source=test.instrumented_behavior_test_source, + instrumented_perf_test_source=test.instrumented_perf_test_source, + behavior_file_path=test.behavior_file_path, + perf_file_path=test.perf_file_path, + ) + ) + return GeneratedTestsList(generated_tests=modified_tests) + + def add_global_declarations(self, optimized_code: str, original_source: str, module_abspath: Path) -> str: + from codeflash.languages.javascript.code_replacer import _add_global_declarations_for_language + + return _add_global_declarations_for_language(optimized_code, original_source, module_abspath, self.language) + + def extract_calling_function_source(self, source_code: str, function_name: str, ref_line: int) -> str | None: + from codeflash.languages.javascript.treesitter import extract_calling_function_source + + return extract_calling_function_source(source_code, function_name, ref_line) + + def _build_runtime_map( + self, inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path + ) -> dict[str, int]: + from codeflash.languages.javascript.edit_tests import resolve_js_test_module_path + + unique_inv_ids: dict[str, int] = {} + for inv_id, runtimes in inv_id_runtimes.items(): + test_qualified_name = ( + inv_id.test_class_name + "." + inv_id.test_function_name # type: ignore[operator] + if inv_id.test_class_name + else inv_id.test_function_name + ) + if not test_qualified_name: + continue + abs_path = resolve_js_test_module_path(inv_id.test_module_path, tests_project_rootdir) + + abs_path_str = str(abs_path.resolve().with_suffix("")) + if "__unit_test_" not in abs_path_str and "__perf_test_" not in abs_path_str: + continue + + key = test_qualified_name + "#" + abs_path_str + parts = inv_id.iteration_id.split("_").__len__() # type: ignore[union-attr] + cur_invid = ( + inv_id.iteration_id.split("_")[0] if parts < 3 else "_".join(inv_id.iteration_id.split("_")[:-1]) + ) # type: ignore[union-attr] + match_key = key + "#" + cur_invid + if match_key not in unique_inv_ids: + unique_inv_ids[match_key] = 0 + unique_inv_ids[match_key] += min(runtimes) + return unique_inv_ids + # === Test Result Comparison === def compare_test_results( @@ -1738,15 +1920,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. @@ -1942,6 +2115,9 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: logger.error("Could not install codeflash. Please run: npm install --save-dev codeflash") return False + def create_dependency_resolver(self, project_root: Path) -> None: + return None + def instrument_existing_test( self, test_path: Path, @@ -2098,6 +2274,10 @@ def run_behavioral_tests( candidate_index=candidate_index, ) + # JavaScript/TypeScript benchmarking uses high max_loops like Python (100,000) + # The actual loop count is limited by target_duration_seconds, not max_loops + JS_BENCHMARKING_MAX_LOOPS = 100_000 + def run_benchmarking_tests( self, test_paths: Any, @@ -2130,10 +2310,15 @@ def run_benchmarking_tests( from codeflash.languages.test_framework import get_js_test_framework_or_default framework = test_framework or get_js_test_framework_or_default() + logger.debug("run_benchmarking_tests called with framework=%s", framework) + + # Use JS-specific high max_loops - actual loop count is limited by target_duration + effective_max_loops = self.JS_BENCHMARKING_MAX_LOOPS if framework == "vitest": from codeflash.languages.javascript.vitest_runner import run_vitest_benchmarking_tests + logger.debug("Dispatching to run_vitest_benchmarking_tests") return run_vitest_benchmarking_tests( test_paths=test_paths, test_env=test_env, @@ -2141,7 +2326,7 @@ def run_benchmarking_tests( timeout=timeout, project_root=project_root, min_loops=min_loops, - max_loops=max_loops, + max_loops=effective_max_loops, target_duration_ms=int(target_duration_seconds * 1000), ) @@ -2154,7 +2339,7 @@ def run_benchmarking_tests( timeout=timeout, project_root=project_root, min_loops=min_loops, - max_loops=max_loops, + max_loops=effective_max_loops, target_duration_ms=int(target_duration_seconds * 1000), ) diff --git a/codeflash/languages/javascript/test_runner.py b/codeflash/languages/javascript/test_runner.py index c65adfa7b..3a193602b 100644 --- a/codeflash/languages/javascript/test_runner.py +++ b/codeflash/languages/javascript/test_runner.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import os import subprocess import time from pathlib import Path @@ -21,6 +22,25 @@ if TYPE_CHECKING: from codeflash.models.models import TestFiles +# Track created config files (jest configs and tsconfigs) for cleanup +_created_config_files: set[Path] = set() + + +def get_created_config_files() -> list[Path]: + """Get list of config files created by codeflash for cleanup. + + Returns: + List of paths to created config files (jest.codeflash.config.js, tsconfig.codeflash.json) + that should be cleaned up after optimization. + + """ + return list(_created_config_files) + + +def clear_created_config_files() -> None: + """Clear the set of tracked config files after cleanup.""" + _created_config_files.clear() + def _detect_bundler_module_resolution(project_root: Path) -> bool: """Detect if the project uses moduleResolution: 'bundler' in tsconfig. @@ -163,6 +183,7 @@ def _create_codeflash_tsconfig(project_root: Path) -> Path: try: codeflash_tsconfig_path.write_text(json.dumps(codeflash_tsconfig, indent=2)) + _created_config_files.add(codeflash_tsconfig_path) logger.debug(f"Created {codeflash_tsconfig_path} with Node moduleResolution") except Exception as e: logger.warning(f"Failed to create codeflash tsconfig: {e}") @@ -170,70 +191,142 @@ def _create_codeflash_tsconfig(project_root: Path) -> Path: return codeflash_tsconfig_path -def _create_codeflash_jest_config(project_root: Path, original_jest_config: Path | None) -> Path | None: - """Create a Jest config that uses the codeflash tsconfig for ts-jest. +def _has_ts_jest_dependency(project_root: Path) -> bool: + """Check if the project has ts-jest as a dependency. + + Args: + project_root: Root of the project. + + Returns: + True if ts-jest is found in dependencies or devDependencies. + + """ + package_json = project_root / "package.json" + if not package_json.exists(): + return False + + try: + content = json.loads(package_json.read_text()) + deps = {**content.get("dependencies", {}), **content.get("devDependencies", {})} + return "ts-jest" in deps + except (json.JSONDecodeError, OSError): + return False + + +def _create_codeflash_jest_config( + project_root: Path, original_jest_config: Path | None, *, for_esm: bool = False +) -> Path | None: + """Create a Jest config that handles ESM packages and TypeScript properly. Args: project_root: Root of the project. original_jest_config: Path to the original Jest config, or None. + for_esm: If True, configure for ESM package transformation. Returns: Path to the codeflash Jest config, or None if creation failed. """ - codeflash_jest_config_path = project_root / "jest.codeflash.config.js" + # For ESM projects (type: module), use .cjs extension since config uses CommonJS require/module.exports + # This prevents "ReferenceError: module is not defined" errors + is_esm = _is_esm_project(project_root) + config_ext = ".cjs" if is_esm else ".js" - # If it already exists, use it + # Create codeflash config in the same directory as the original config + # This ensures relative paths work correctly + if original_jest_config: + codeflash_jest_config_path = original_jest_config.parent / f"jest.codeflash.config{config_ext}" + else: + codeflash_jest_config_path = project_root / f"jest.codeflash.config{config_ext}" + + # If it already exists, use it (check both extensions) if codeflash_jest_config_path.exists(): logger.debug(f"Using existing {codeflash_jest_config_path}") return codeflash_jest_config_path - # Create a wrapper Jest config that uses tsconfig.codeflash.json + # Also check if the alternate extension exists + alt_ext = ".js" if is_esm else ".cjs" + alt_path = codeflash_jest_config_path.with_suffix(alt_ext) + if alt_path.exists(): + logger.debug(f"Using existing {alt_path}") + return alt_path + + # Common ESM-only packages that need to be transformed + # These packages ship only ESM and will cause "Cannot use import statement" errors + esm_packages = [ + "p-queue", + "p-limit", + "p-timeout", + "yocto-queue", + "eventemitter3", + "chalk", + "ora", + "strip-ansi", + "ansi-regex", + "string-width", + "wrap-ansi", + "is-unicode-supported", + "is-interactive", + "log-symbols", + "figures", + ] + esm_pattern = "|".join(esm_packages) + + # Check if ts-jest is available in the project + has_ts_jest = _has_ts_jest_dependency(project_root) + + # Build transform config only if ts-jest is available + if has_ts_jest: + transform_config = """ + // Ensure TypeScript files are transformed using ts-jest + transform: { + '^.+\\\\.(ts|tsx)$': ['ts-jest', { isolatedModules: true }], + // Use ts-jest for JS files in ESM packages too + '^.+\\\\.js$': ['ts-jest', { isolatedModules: true }], + },""" + else: + transform_config = "" + logger.debug("ts-jest not found in project dependencies, skipping transform config") + + # Create a wrapper Jest config if original_jest_config: - # Extend the original config - jest_config_content = f"""// Auto-generated by codeflash for bundler moduleResolution compatibility -const originalConfig = require('./{original_jest_config.name}'); + # Since codeflash config is in the same directory as original, use simple relative path + config_require_path = f"./{original_jest_config.name}" -const tsJestOptions = {{ - isolatedModules: true, - tsconfig: 'tsconfig.codeflash.json', -}}; + # Extend the original config + jest_config_content = f"""// Auto-generated by codeflash for ESM compatibility +const originalConfig = require('{config_require_path}'); module.exports = {{ ...originalConfig, - transform: {{ - ...originalConfig.transform, - '^.+\\\\.tsx?$': ['ts-jest', tsJestOptions], - }}, - globals: {{ - ...originalConfig.globals, - 'ts-jest': tsJestOptions, - }}, + // Transform ESM packages that don't work with Jest's default config + // Pattern handles both npm/yarn (node_modules/pkg) and pnpm (node_modules/.pnpm/pkg@version/node_modules/pkg) + transformIgnorePatterns: [ + 'node_modules/(?!(\\\\.pnpm/)?({esm_pattern}))', + ],{transform_config} }}; """ else: - # Create a minimal Jest config for TypeScript - jest_config_content = """// Auto-generated by codeflash for bundler moduleResolution compatibility -const tsJestOptions = { - isolatedModules: true, - tsconfig: 'tsconfig.codeflash.json', -}; - -module.exports = { + # Create a minimal Jest config for TypeScript with ESM support + jest_config_content = f"""// Auto-generated by codeflash for ESM compatibility +module.exports = {{ verbose: true, testEnvironment: 'node', testRegex: '\\\\.(test|spec)\\\\.(js|ts|tsx)$', - testPathIgnorePatterns: ['/dist/', '/node_modules/'], - transform: { - '^.+\\\\.tsx?$': ['ts-jest', tsJestOptions], - }, + testPathIgnorePatterns: ['/dist/'], + // Transform ESM packages that don't work with Jest's default config + // Pattern handles both npm/yarn and pnpm directory structures + transformIgnorePatterns: [ + 'node_modules/(?!(\\\\.pnpm/)?({esm_pattern}))', + ],{transform_config} moduleFileExtensions: ['ts', 'tsx', 'js', 'jsx', 'json', 'node'], -}; +}}; """ try: codeflash_jest_config_path.write_text(jest_config_content) - logger.debug(f"Created {codeflash_jest_config_path} with codeflash tsconfig") + _created_config_files.add(codeflash_jest_config_path) + logger.debug(f"Created {codeflash_jest_config_path} with ESM package support") return codeflash_jest_config_path except Exception as e: logger.warning(f"Failed to create codeflash Jest config: {e}") @@ -323,6 +416,55 @@ def _find_monorepo_root(start_path: Path) -> Path | None: return None +def _get_jest_major_version(project_root: Path) -> int | None: + """Detect the major version of Jest installed in the project. + + Args: + project_root: Root of the project to check. + + Returns: + Major version number (e.g., 29, 30), or None if not detected. + + """ + # First try to check package.json for explicit version + package_json = project_root / "package.json" + if package_json.exists(): + try: + content = json.loads(package_json.read_text()) + deps = {**content.get("devDependencies", {}), **content.get("dependencies", {})} + jest_version = deps.get("jest", "") + # Parse version like "30.0.5", "^30.0.5", "~30.0.5" + if jest_version: + # Strip leading version prefixes (^, ~, =, v) + version_str = jest_version.lstrip("^~=v") + if version_str and version_str[0].isdigit(): + major = version_str.split(".")[0] + if major.isdigit(): + return int(major) + except (json.JSONDecodeError, OSError): + pass + + # Also check monorepo root + monorepo_root = _find_monorepo_root(project_root) + if monorepo_root and monorepo_root != project_root: + monorepo_package = monorepo_root / "package.json" + if monorepo_package.exists(): + try: + content = json.loads(monorepo_package.read_text()) + deps = {**content.get("devDependencies", {}), **content.get("dependencies", {})} + jest_version = deps.get("jest", "") + if jest_version: + version_str = jest_version.lstrip("^~=v") + if version_str and version_str[0].isdigit(): + major = version_str.split(".")[0] + if major.isdigit(): + return int(major) + except (json.JSONDecodeError, OSError): + pass + + return None + + def _find_jest_config(project_root: Path) -> Path | None: """Find Jest configuration file in the project. @@ -535,7 +677,6 @@ def run_jest_behavioral_tests( # Get test files to run test_files = [str(file.instrumented_behavior_file_path) for file in test_paths.test_files] - # Use provided project_root, or detect it as fallback if project_root is None and test_files: first_test_file = Path(test_files[0]) @@ -610,13 +751,25 @@ def run_jest_behavioral_tests( # Configure ESM support if project uses ES Modules _configure_esm_environment(jest_env, effective_cwd) + # Increase Node.js heap size for large TypeScript projects + # Default heap is often not enough for monorepos with many dependencies + existing_node_options = jest_env.get("NODE_OPTIONS", "") + if "--max-old-space-size" not in existing_node_options: + jest_env["NODE_OPTIONS"] = f"{existing_node_options} --max-old-space-size=4096".strip() + logger.debug(f"Running Jest tests with command: {' '.join(jest_cmd)}") + # Calculate subprocess timeout: needs to be much larger than per-test timeout + # to account for Jest startup, TypeScript compilation, module loading, etc. + # Use at least 120 seconds, or 10x the per-test timeout, whichever is larger + subprocess_timeout = max(120, (timeout or 15) * 10, 600) if timeout else 600 + start_time_ns = time.perf_counter_ns() try: run_args = get_cross_platform_subprocess_run_args( - cwd=effective_cwd, env=jest_env, timeout=timeout or 600, check=False, text=True, capture_output=True + cwd=effective_cwd, env=jest_env, timeout=subprocess_timeout, check=False, text=True, capture_output=True ) + logger.debug(f"Jest subprocess timeout: {subprocess_timeout}s (per-test timeout: {timeout}s)") result = subprocess.run(jest_cmd, **run_args) # noqa: PLW1510 # Jest sends console.log output to stderr by default - move it to stdout # so our timing markers (printed via console.log) are in the expected place @@ -634,12 +787,12 @@ def run_jest_behavioral_tests( # This helps debug issues like import errors that cause Jest to fail early if result.returncode != 0 and not result_file_path.exists(): logger.warning( - f"Jest failed with returncode={result.returncode} and no XML output created.\n" + f"Jest failed with returncode={result.returncode}.\n" f"Jest stdout: {result.stdout[:2000] if result.stdout else '(empty)'}\n" f"Jest stderr: {result.stderr[:500] if result.stderr else '(empty)'}" ) except subprocess.TimeoutExpired: - logger.warning(f"Jest tests timed out after {timeout}s") + logger.warning(f"Jest tests timed out after {subprocess_timeout}s") result = subprocess.CompletedProcess(args=jest_cmd, returncode=-1, stdout="", stderr="Test execution timed out") except FileNotFoundError: logger.error("Jest not found. Make sure Jest is installed (npm install jest)") @@ -774,25 +927,28 @@ def run_jest_benchmarking_tests( # Get performance test files test_files = [str(file.benchmarking_file_path) for file in test_paths.test_files if file.benchmarking_file_path] - # Use provided project_root, or detect it as fallback if project_root is None and test_files: first_test_file = Path(test_files[0]) project_root = _find_node_project_root(first_test_file) effective_cwd = project_root if project_root else cwd - logger.debug(f"Jest benchmarking working directory: {effective_cwd}") # Ensure the codeflash npm package is installed _ensure_runtime_files(effective_cwd) - # Build Jest command for performance tests with custom loop runner + # Detect Jest version for logging + jest_major_version = _get_jest_major_version(effective_cwd) + if jest_major_version: + logger.debug(f"Jest {jest_major_version} detected - using loop-runner for batched looping") + + # Build Jest command for performance tests jest_cmd = [ "npx", "jest", "--reporters=default", "--reporters=jest-junit", - "--runInBand", # Ensure serial execution even though runner enforces it + "--runInBand", # Ensure serial execution "--forceExit", "--runner=codeflash/loop-runner", # Use custom loop runner for in-process looping ] @@ -844,9 +1000,25 @@ def run_jest_benchmarking_tests( jest_env["CODEFLASH_PERF_STABILITY_CHECK"] = "true" if stability_check else "false" jest_env["CODEFLASH_LOOP_INDEX"] = "1" # Initial value for compatibility + # Enable console output for timing markers + # Some projects mock console.log in test setup (e.g., based on LOG_LEVEL or DEBUG) + # We need console.log to work for capturePerf timing markers + jest_env["LOG_LEVEL"] = "info" # Disable console.log mocking in projects that check LOG_LEVEL + jest_env["DEBUG"] = "1" # Disable console.log mocking in projects that check DEBUG + + # Debug logging for loop behavior verification (set CODEFLASH_DEBUG_LOOPS=true to enable) + if os.environ.get("CODEFLASH_DEBUG_LOOPS") == "true": + jest_env["CODEFLASH_DEBUG_LOOPS"] = "true" + logger.info("Loop debug logging enabled - will show capturePerf loop details") + # Configure ESM support if project uses ES Modules _configure_esm_environment(jest_env, effective_cwd) + # Increase Node.js heap size for large TypeScript projects + existing_node_options = jest_env.get("NODE_OPTIONS", "") + if "--max-old-space-size" not in existing_node_options: + jest_env["NODE_OPTIONS"] = f"{existing_node_options} --max-old-space-size=4096".strip() + # Total timeout for the entire benchmark run (longer than single-loop timeout) # Account for startup overhead + target duration + buffer total_timeout = max(120, (target_duration_ms // 1000) + 60, timeout or 120) @@ -872,6 +1044,10 @@ def run_jest_benchmarking_tests( # Create result with combined stdout result = subprocess.CompletedProcess(args=result.args, returncode=result.returncode, stdout=stdout, stderr="") + if result.returncode != 0: + logger.info(f"Jest benchmarking failed with return code {result.returncode}") + logger.info(f"Jest benchmarking stdout: {result.stdout}") + logger.info(f"Jest benchmarking stderr: {result.stderr}") except subprocess.TimeoutExpired: logger.warning(f"Jest benchmarking timed out after {total_timeout}s") @@ -882,7 +1058,6 @@ def run_jest_benchmarking_tests( wall_clock_seconds = time.time() - total_start_time logger.debug(f"Jest benchmarking completed in {wall_clock_seconds:.2f}s") - return result_file_path, result @@ -985,6 +1160,11 @@ def run_jest_line_profile_tests( # Configure ESM support if project uses ES Modules _configure_esm_environment(jest_env, effective_cwd) + # Increase Node.js heap size for large TypeScript projects + existing_node_options = jest_env.get("NODE_OPTIONS", "") + if "--max-old-space-size" not in existing_node_options: + jest_env["NODE_OPTIONS"] = f"{existing_node_options} --max-old-space-size=4096".strip() + subprocess_timeout = timeout or 600 logger.debug(f"Running Jest line profile tests: {' '.join(jest_cmd)}") diff --git a/codeflash/languages/javascript/treesitter.py b/codeflash/languages/javascript/treesitter.py new file mode 100644 index 000000000..c00cb228e --- /dev/null +++ b/codeflash/languages/javascript/treesitter.py @@ -0,0 +1,1826 @@ +"""Tree-sitter utilities for cross-language code analysis. + +This module provides a unified interface for parsing and analyzing code +across multiple languages using tree-sitter. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +from tree_sitter import Language, Parser + +if TYPE_CHECKING: + from pathlib import Path + + from tree_sitter import Node, Tree + +logger = logging.getLogger(__name__) + + +class TreeSitterLanguage(Enum): + """Supported tree-sitter languages.""" + + JAVASCRIPT = "javascript" + TYPESCRIPT = "typescript" + TSX = "tsx" + + +# Lazy-loaded language instances +_LANGUAGE_CACHE: dict[TreeSitterLanguage, Language] = {} + + +def _get_language(lang: TreeSitterLanguage) -> Language: + """Get a tree-sitter Language instance, with lazy loading.""" + if lang not in _LANGUAGE_CACHE: + if lang == TreeSitterLanguage.JAVASCRIPT: + import tree_sitter_javascript + + _LANGUAGE_CACHE[lang] = Language(tree_sitter_javascript.language()) + elif lang == TreeSitterLanguage.TYPESCRIPT: + import tree_sitter_typescript + + _LANGUAGE_CACHE[lang] = Language(tree_sitter_typescript.language_typescript()) + elif lang == TreeSitterLanguage.TSX: + import tree_sitter_typescript + + _LANGUAGE_CACHE[lang] = Language(tree_sitter_typescript.language_tsx()) + return _LANGUAGE_CACHE[lang] + + +@dataclass +class FunctionNode: + """Represents a function found by tree-sitter analysis.""" + + name: str + node: Node + start_line: int + end_line: int + start_col: int + end_col: int + is_async: bool + is_method: bool + is_arrow: bool + is_generator: bool + class_name: str | None + parent_function: str | None + source_text: str + doc_start_line: int | None = None # Line where JSDoc comment starts (or None if no JSDoc) + is_exported: bool = False # Whether the function is exported + + +@dataclass +class ImportInfo: + """Represents an import statement.""" + + module_path: str # The path being imported from + default_import: str | None # Default import name (import X from ...) + named_imports: list[tuple[str, str | None]] # [(name, alias), ...] + namespace_import: str | None # Namespace import (import * as X from ...) + is_type_only: bool # TypeScript type-only import + start_line: int + end_line: int + + +@dataclass +class ExportInfo: + """Represents an export statement.""" + + exported_names: list[tuple[str, str | None]] # [(name, alias), ...] for named exports + default_export: str | None # Name of default exported function/class/value + is_reexport: bool # Whether this is a re-export (export { x } from './other') + reexport_source: str | None # Module path for re-exports + start_line: int + end_line: int + # Functions passed as arguments to wrapper calls in default exports + # e.g., export default curry(traverseEntity) -> ["traverseEntity"] + wrapped_default_args: list[str] | None = None + + +@dataclass +class ModuleLevelDeclaration: + """Represents a module-level (global) variable or constant declaration.""" + + name: str # Variable/constant name + declaration_type: str # "const", "let", "var", "class", "enum", "type", "interface" + source_code: str # Full declaration source code + start_line: int + end_line: int + is_exported: bool # Whether the declaration is exported + + +@dataclass +class TypeDefinition: + """Represents a type definition (interface, type alias, class, or enum).""" + + name: str # Type name + definition_type: str # "interface", "type", "class", "enum" + source_code: str # Full definition source code + start_line: int + end_line: int + is_exported: bool # Whether the definition is exported + file_path: Path | None = None # File where the type is defined + + +class TreeSitterAnalyzer: + """Cross-language code analysis using tree-sitter. + + This class provides methods to parse and analyze JavaScript/TypeScript code, + finding functions, imports, and other code structures. + """ + + def __init__(self, language: TreeSitterLanguage | str) -> None: + """Initialize the analyzer for a specific language. + + Args: + language: The language to analyze (TreeSitterLanguage enum or string). + + """ + if isinstance(language, str): + language = TreeSitterLanguage(language) + self.language = language + self._parser: Parser | None = None + + @property + def parser(self) -> Parser: + """Get the parser, creating it lazily.""" + if self._parser is None: + self._parser = Parser(_get_language(self.language)) + return self._parser + + def parse(self, source: str | bytes) -> Tree: + """Parse source code into a tree-sitter tree. + + Args: + source: Source code as string or bytes. + + Returns: + The parsed tree. + + """ + if isinstance(source, str): + source = source.encode("utf8") + return self.parser.parse(source) + + def get_node_text(self, node: Node, source: bytes) -> str: + """Extract the source text for a tree-sitter node. + + Args: + node: The tree-sitter node. + source: The source code as bytes. + + Returns: + The text content of the node. + + """ + return source[node.start_byte : node.end_byte].decode("utf8") + + def find_functions( + self, source: str, include_methods: bool = True, include_arrow_functions: bool = True, require_name: bool = True + ) -> list[FunctionNode]: + """Find all function definitions in source code. + + Args: + source: The source code to analyze. + include_methods: Whether to include class methods. + include_arrow_functions: Whether to include arrow functions. + require_name: Whether to require functions to have names. + + Returns: + List of FunctionNode objects describing found functions. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + functions: list[FunctionNode] = [] + + self._walk_tree_for_functions( + tree.root_node, + source_bytes, + functions, + include_methods=include_methods, + include_arrow_functions=include_arrow_functions, + require_name=require_name, + current_class=None, + current_function=None, + ) + + return functions + + def _walk_tree_for_functions( + self, + node: Node, + source_bytes: bytes, + functions: list[FunctionNode], + include_methods: bool, + include_arrow_functions: bool, + require_name: bool, + current_class: str | None, + current_function: str | None, + ) -> None: + """Recursively walk the tree to find function definitions.""" + # Function types in JavaScript/TypeScript + function_types = { + "function_declaration", + "function_expression", + "generator_function_declaration", + "generator_function", + } + + if include_arrow_functions: + function_types.add("arrow_function") + + if include_methods: + function_types.add("method_definition") + + # Track class context + new_class = current_class + new_function = current_function + + if node.type in {"class_declaration", "class"}: + # Get class name + name_node = node.child_by_field_name("name") + if name_node: + new_class = self.get_node_text(name_node, source_bytes) + + if node.type in function_types: + func_info = self._extract_function_info(node, source_bytes, current_class, current_function) + + if func_info: + # Check if we should include this function + should_include = True + + if require_name and not func_info.name: + should_include = False + + if func_info.is_method and not include_methods: + should_include = False + + if func_info.is_arrow and not include_arrow_functions: + should_include = False + + # Skip arrow functions that are object properties (e.g., { foo: () => {} }) + # These are not standalone functions - they're values in object literals + if func_info.is_arrow and node.parent and node.parent.type == "pair": + should_include = False + + if should_include: + functions.append(func_info) + + # Track as current function for nested functions + if func_info.name: + new_function = func_info.name + + # Recurse into children + for child in node.children: + self._walk_tree_for_functions( + child, + source_bytes, + functions, + include_methods=include_methods, + include_arrow_functions=include_arrow_functions, + require_name=require_name, + current_class=new_class, + current_function=new_function if node.type in function_types else current_function, + ) + + def _extract_function_info( + self, node: Node, source_bytes: bytes, current_class: str | None, current_function: str | None + ) -> FunctionNode | None: + """Extract function information from a tree-sitter node.""" + name = "" + is_async = False + is_generator = False + is_method = False + is_arrow = node.type == "arrow_function" + is_exported = False + + # Check for async modifier + for child in node.children: + if child.type == "async": + is_async = True + break + + # Check for generator + if "generator" in node.type: + is_generator = True + + # Check if function is exported + # For function_declaration: check if parent is export_statement + # For arrow functions: check if parent variable_declarator's grandparent is export_statement + # For CommonJS: check module.exports = { name } or exports.name = ... + is_exported = self._is_node_exported(node, source_bytes) + + # Get function name based on node type + if node.type in ("function_declaration", "generator_function_declaration"): + name_node = node.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + else: + # Fallback: search for identifier child (some tree-sitter versions) + for child in node.children: + if child.type == "identifier": + name = self.get_node_text(child, source_bytes) + break + elif node.type == "method_definition": + is_method = True + name_node = node.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + elif node.type in ("function_expression", "generator_function"): + # Check if assigned to a variable + name_node = node.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + else: + # Try to get name from parent assignment + name = self._get_name_from_assignment(node, source_bytes) + elif node.type == "arrow_function": + # Arrow functions get names from variable declarations + name = self._get_name_from_assignment(node, source_bytes) + + # Get source text + source_text = self.get_node_text(node, source_bytes) + + # Find preceding JSDoc comment + doc_start_line = self._find_preceding_jsdoc(node, source_bytes) + + return FunctionNode( + name=name, + node=node, + start_line=node.start_point[0] + 1, # Convert to 1-indexed + end_line=node.end_point[0] + 1, + start_col=node.start_point[1], + end_col=node.end_point[1], + is_async=is_async, + is_method=is_method, + is_arrow=is_arrow, + is_generator=is_generator, + class_name=current_class if is_method else None, + parent_function=current_function, + source_text=source_text, + doc_start_line=doc_start_line, + is_exported=is_exported, + ) + + def _is_node_exported(self, node: Node, source_bytes: bytes | None = None) -> bool: + """Check if a function node is exported. + + Handles various export patterns: + - export function foo() {} + - export const foo = () => {} + - export default function foo() {} + - Class methods in exported classes + - module.exports = { foo } (CommonJS) + - exports.foo = ... (CommonJS) + + Args: + node: The function node to check. + source_bytes: Source code bytes (needed for CommonJS export detection). + + Returns: + True if the function is exported, False otherwise. + + """ + # Check direct parent for export_statement + if node.parent and node.parent.type == "export_statement": + return True + + # For arrow functions and function expressions assigned to variables + # e.g., export const foo = () => {} + if node.type in ("arrow_function", "function_expression", "generator_function"): + parent = node.parent + if parent and parent.type == "variable_declarator": + grandparent = parent.parent + if grandparent and grandparent.type in ("lexical_declaration", "variable_declaration"): + great_grandparent = grandparent.parent + if great_grandparent and great_grandparent.type == "export_statement": + return True + + # For methods in exported classes + if node.type == "method_definition": + # Walk up to find class_declaration + current = node.parent + while current: + if current.type in ("class_declaration", "class"): + # Check if this class is exported via ES module export + if current.parent and current.parent.type == "export_statement": + return True + # Check if class is exported via CommonJS + if source_bytes: + class_name_node = current.child_by_field_name("name") + if class_name_node: + class_name = self.get_node_text(class_name_node, source_bytes) + if self._is_name_in_commonjs_exports(node, class_name, source_bytes): + return True + break + current = current.parent + + # Check CommonJS exports: module.exports = { foo } or exports.foo = ... + if source_bytes: + func_name = self._get_function_name_for_export_check(node, source_bytes) + if func_name and self._is_name_in_commonjs_exports(node, func_name, source_bytes): + return True + + return False + + def _get_function_name_for_export_check(self, node: Node, source_bytes: bytes) -> str | None: + """Get the function name for export checking.""" + if node.type in ("function_declaration", "generator_function_declaration"): + name_node = node.child_by_field_name("name") + if name_node: + return self.get_node_text(name_node, source_bytes) + elif node.type in ("arrow_function", "function_expression", "generator_function"): + # Get name from variable assignment + parent = node.parent + if parent and parent.type == "variable_declarator": + name_node = parent.child_by_field_name("name") + if name_node and name_node.type == "identifier": + return self.get_node_text(name_node, source_bytes) + return None + + def _is_name_in_commonjs_exports(self, node: Node, name: str, source_bytes: bytes) -> bool: + """Check if a name is exported via CommonJS module.exports or exports. + + Handles patterns like: + - module.exports = { foo, bar } + - module.exports = { foo: someFunc } + - exports.foo = ... + - module.exports.foo = ... + + Args: + node: Any node in the tree (used to find the program root). + name: The name to check for in exports. + source_bytes: Source code bytes. + + Returns: + True if the name is in CommonJS exports. + + """ + # Walk up to find program root + root = node + while root.parent: + root = root.parent + + # Search for CommonJS export patterns in program children + for child in root.children: + if child.type == "expression_statement": + # Look for assignment expressions + for expr in child.children: + if expr.type == "assignment_expression": + if self._check_commonjs_assignment_exports(expr, name, source_bytes): + return True + + return False + + def _check_commonjs_assignment_exports(self, node: Node, name: str, source_bytes: bytes) -> bool: + """Check if a CommonJS assignment exports the given name.""" + left_node = node.child_by_field_name("left") + right_node = node.child_by_field_name("right") + + if not left_node or not right_node: + return False + + left_text = self.get_node_text(left_node, source_bytes) + + # Check module.exports = { name, ... } or module.exports = { key: name, ... } + if left_text == "module.exports" and right_node.type == "object": + for child in right_node.children: + if child.type == "shorthand_property_identifier": + # { foo } - shorthand export + if self.get_node_text(child, source_bytes) == name: + return True + elif child.type == "pair": + # { key: value } - check both key and value + key_node = child.child_by_field_name("key") + value_node = child.child_by_field_name("value") + if key_node and self.get_node_text(key_node, source_bytes) == name: + return True + if value_node and value_node.type == "identifier": + if self.get_node_text(value_node, source_bytes) == name: + return True + + # Check module.exports = name (single export) + if left_text == "module.exports" and right_node.type == "identifier": + if self.get_node_text(right_node, source_bytes) == name: + return True + + # Check module.exports.name = ... or exports.name = ... + if left_text in {f"module.exports.{name}", f"exports.{name}"}: + return True + + return False + + def _find_preceding_jsdoc(self, node: Node, source_bytes: bytes) -> int | None: + """Find JSDoc comment immediately preceding a function node. + + For regular functions, looks at the previous sibling of the function node. + For arrow functions assigned to variables, looks at the previous sibling + of the variable declaration. + + Args: + node: The function node to find JSDoc for. + source_bytes: The source code as bytes. + + Returns: + The start line (1-indexed) of the JSDoc, or None if no JSDoc found. + + """ + target_node = node + + # For arrow functions, look at parent variable declaration + if node.type == "arrow_function": + parent = node.parent + if parent and parent.type == "variable_declarator": + grandparent = parent.parent + if grandparent and grandparent.type in ("lexical_declaration", "variable_declaration"): + target_node = grandparent + + # For function expressions assigned to variables, also look at parent + if node.type in ("function_expression", "generator_function"): + parent = node.parent + if parent and parent.type == "variable_declarator": + grandparent = parent.parent + if grandparent and grandparent.type in ("lexical_declaration", "variable_declaration"): + target_node = grandparent + + # Get the previous sibling node + prev_sibling = target_node.prev_named_sibling + + # Check if it's a comment node with JSDoc pattern + if prev_sibling and prev_sibling.type == "comment": + comment_text = self.get_node_text(prev_sibling, source_bytes) + if comment_text.strip().startswith("/**"): + # Verify it's immediately preceding (no blank lines between) + comment_end_line = prev_sibling.end_point[0] + function_start_line = target_node.start_point[0] + if function_start_line - comment_end_line <= 1: + return prev_sibling.start_point[0] + 1 # 1-indexed + + return None + + def _get_name_from_assignment(self, node: Node, source_bytes: bytes) -> str: + """Try to extract function name from parent variable declaration or assignment. + + Handles patterns like: + - const foo = () => {} + - const foo = function() {} + - let bar = function() {} + - obj.method = () => {} + """ + parent = node.parent + if parent is None: + return "" + + # Check for variable declarator: const foo = ... + if parent.type == "variable_declarator": + name_node = parent.child_by_field_name("name") + if name_node: + return self.get_node_text(name_node, source_bytes) + + # Check for assignment expression: foo = ... + if parent.type == "assignment_expression": + left_node = parent.child_by_field_name("left") + if left_node: + if left_node.type == "identifier": + return self.get_node_text(left_node, source_bytes) + if left_node.type == "member_expression": + # For obj.method = ..., get the property name + prop_node = left_node.child_by_field_name("property") + if prop_node: + return self.get_node_text(prop_node, source_bytes) + + # Check for property in object: { foo: () => {} } + if parent.type == "pair": + key_node = parent.child_by_field_name("key") + if key_node: + return self.get_node_text(key_node, source_bytes) + + return "" + + def find_imports(self, source: str) -> list[ImportInfo]: + """Find all import statements in source code. + + Args: + source: The source code to analyze. + + Returns: + List of ImportInfo objects describing imports. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + imports: list[ImportInfo] = [] + + self._walk_tree_for_imports(tree.root_node, source_bytes, imports) + + return imports + + def _walk_tree_for_imports( + self, node: Node, source_bytes: bytes, imports: list[ImportInfo], in_function: bool = False + ) -> None: + """Recursively walk the tree to find import statements. + + Args: + node: Current node to check. + source_bytes: Source code bytes. + imports: List to append found imports to. + in_function: Whether we're currently inside a function/method body. + + """ + # Track when we enter function/method bodies + # These node types contain function/method bodies where require() should not be treated as imports + function_body_types = { + "function_declaration", + "method_definition", + "arrow_function", + "function_expression", + "function", # Generic function in some grammars + } + + if node.type == "import_statement": + import_info = self._extract_import_info(node, source_bytes) + if import_info: + imports.append(import_info) + + # Also handle require() calls for CommonJS, but only at module level + # require() inside functions is a dynamic import, not a module import + if node.type == "call_expression" and not in_function: + func_node = node.child_by_field_name("function") + if func_node and self.get_node_text(func_node, source_bytes) == "require": + import_info = self._extract_require_info(node, source_bytes) + if import_info: + imports.append(import_info) + + # Update in_function flag for children + child_in_function = in_function or node.type in function_body_types + + for child in node.children: + self._walk_tree_for_imports(child, source_bytes, imports, child_in_function) + + def _extract_import_info(self, node: Node, source_bytes: bytes) -> ImportInfo | None: + """Extract import information from an import statement node.""" + module_path = "" + default_import = None + named_imports: list[tuple[str, str | None]] = [] + namespace_import = None + is_type_only = False + + # Get the module path (source) + source_node = node.child_by_field_name("source") + if source_node: + # Remove quotes from string + module_path = self.get_node_text(source_node, source_bytes).strip("'\"") + + # Check for type-only import (TypeScript) + for child in node.children: + if child.type == "type" or self.get_node_text(child, source_bytes) == "type": + is_type_only = True + break + + # Process import clause + for child in node.children: + if child.type == "import_clause": + self._process_import_clause(child, source_bytes, default_import, named_imports, namespace_import) + # Re-extract after processing + for clause_child in child.children: + if clause_child.type == "identifier": + default_import = self.get_node_text(clause_child, source_bytes) + elif clause_child.type == "named_imports": + for spec in clause_child.children: + if spec.type == "import_specifier": + name_node = spec.child_by_field_name("name") + alias_node = spec.child_by_field_name("alias") + if name_node: + name = self.get_node_text(name_node, source_bytes) + alias = self.get_node_text(alias_node, source_bytes) if alias_node else None + named_imports.append((name, alias)) + elif clause_child.type == "namespace_import": + # import * as X + for ns_child in clause_child.children: + if ns_child.type == "identifier": + namespace_import = self.get_node_text(ns_child, source_bytes) + + if not module_path: + return None + + return ImportInfo( + module_path=module_path, + default_import=default_import, + named_imports=named_imports, + namespace_import=namespace_import, + is_type_only=is_type_only, + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + ) + + def _process_import_clause( + self, + node: Node, + source_bytes: bytes, + default_import: str | None, + named_imports: list[tuple[str, str | None]], + namespace_import: str | None, + ) -> None: + """Process an import clause to extract imports.""" + # This is a helper that modifies the lists in place + # Processing is done inline in _extract_import_info + + def _extract_require_info(self, node: Node, source_bytes: bytes) -> ImportInfo | None: + """Extract import information from a require() call. + + Handles various CommonJS require patterns: + - const foo = require('./module') -> default import + - const { a, b } = require('./module') -> named imports + - const { a: aliasA } = require('./module') -> named imports with alias + - const foo = require('./module').bar -> property access (named import) + - require('./module') -> side effect import + """ + # Handle require().property pattern - the call_expression is inside member_expression + actual_require_node = node + property_access = None + + # Check if this require is part of a member_expression like require('./m').foo + if node.parent and node.parent.type == "member_expression": + member_node = node.parent + prop_node = member_node.child_by_field_name("property") + if prop_node: + property_access = self.get_node_text(prop_node, source_bytes) + # Use the member expression's parent for variable assignment lookup + node = member_node + + args_node = actual_require_node.child_by_field_name("arguments") + if not args_node: + return None + + # Get the first argument (module path) + module_path = "" + for child in args_node.children: + if child.type == "string": + module_path = self.get_node_text(child, source_bytes).strip("'\"") + break + + if not module_path: + return None + + # Try to get the variable name from assignment + default_import = None + named_imports: list[tuple[str, str | None]] = [] + + parent = node.parent + if parent and parent.type == "variable_declarator": + name_node = parent.child_by_field_name("name") + if name_node: + if name_node.type == "identifier": + var_name = self.get_node_text(name_node, source_bytes) + if property_access: + # const foo = require('./module').bar + # This imports 'bar' from the module and assigns to 'foo' + named_imports.append((property_access, var_name if var_name != property_access else None)) + else: + # const foo = require('./module') + default_import = var_name + elif name_node.type == "object_pattern": + # Destructuring: const { a, b } = require('...') + named_imports = self._extract_object_pattern_names(name_node, source_bytes) + elif property_access: + # require('./module').foo without assignment - still track the property access + named_imports.append((property_access, None)) + + return ImportInfo( + module_path=module_path, + default_import=default_import, + named_imports=named_imports, + namespace_import=None, + is_type_only=False, + start_line=actual_require_node.start_point[0] + 1, + end_line=actual_require_node.end_point[0] + 1, + ) + + def _extract_object_pattern_names(self, node: Node, source_bytes: bytes) -> list[tuple[str, str | None]]: + """Extract names from an object pattern (destructuring). + + Handles patterns like: + - { a, b } -> [('a', None), ('b', None)] + - { a: aliasA } -> [('a', 'aliasA')] + - { a, b: aliasB } -> [('a', None), ('b', 'aliasB')] + """ + names: list[tuple[str, str | None]] = [] + + for child in node.children: + if child.type == "shorthand_property_identifier_pattern": + # { a } - shorthand, name equals value + name = self.get_node_text(child, source_bytes) + names.append((name, None)) + elif child.type == "pair_pattern": + # { a: aliasA } - renamed import + key_node = child.child_by_field_name("key") + value_node = child.child_by_field_name("value") + if key_node and value_node: + original_name = self.get_node_text(key_node, source_bytes) + alias = self.get_node_text(value_node, source_bytes) + names.append((original_name, alias)) + + return names + + def find_exports(self, source: str) -> list[ExportInfo]: + """Find all export statements in source code. + + Args: + source: The source code to analyze. + + Returns: + List of ExportInfo objects describing exports. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + exports: list[ExportInfo] = [] + + self._walk_tree_for_exports(tree.root_node, source_bytes, exports) + + return exports + + def _walk_tree_for_exports(self, node: Node, source_bytes: bytes, exports: list[ExportInfo]) -> None: + """Recursively walk the tree to find export statements.""" + # Handle ES module export statements + if node.type == "export_statement": + export_info = self._extract_export_info(node, source_bytes) + if export_info: + exports.append(export_info) + + # Handle CommonJS exports: module.exports = ... or exports.foo = ... + if node.type == "assignment_expression": + export_info = self._extract_commonjs_export(node, source_bytes) + if export_info: + exports.append(export_info) + + for child in node.children: + self._walk_tree_for_exports(child, source_bytes, exports) + + def _extract_export_info(self, node: Node, source_bytes: bytes) -> ExportInfo | None: + """Extract export information from an export statement node.""" + exported_names: list[tuple[str, str | None]] = [] + default_export: str | None = None + is_reexport = False + reexport_source: str | None = None + wrapped_default_args: list[str] | None = None + + # Check for re-export source (export { x } from './other') + source_node = node.child_by_field_name("source") + if source_node: + is_reexport = True + reexport_source = self.get_node_text(source_node, source_bytes).strip("'\"") + + for child in node.children: + # Handle 'export default' + if child.type == "default": + # Find what's being exported as default + for sibling in node.children: + if sibling.type in {"function_declaration", "class_declaration"}: + name_node = sibling.child_by_field_name("name") + default_export = self.get_node_text(name_node, source_bytes) if name_node else "default" + elif sibling.type == "identifier": + default_export = self.get_node_text(sibling, source_bytes) + elif sibling.type in ("arrow_function", "function_expression", "object", "array"): + default_export = "default" + elif sibling.type == "call_expression": + # Handle wrapped exports: export default curry(traverseEntity) + # The default export is the result of the call, but we track + # the wrapped function names for export checking + default_export = "default" + wrapped_default_args = self._extract_call_expression_identifiers(sibling, source_bytes) + break + + # Handle named exports: export { a, b as c } + if child.type == "export_clause": + for spec in child.children: + if spec.type == "export_specifier": + name_node = spec.child_by_field_name("name") + alias_node = spec.child_by_field_name("alias") + if name_node: + name = self.get_node_text(name_node, source_bytes) + alias = self.get_node_text(alias_node, source_bytes) if alias_node else None + exported_names.append((name, alias)) + + # Handle direct exports: export function foo() {} + if child.type == "function_declaration": + name_node = child.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + exported_names.append((name, None)) + + # Handle direct class exports: export class Foo {} + if child.type == "class_declaration": + name_node = child.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + exported_names.append((name, None)) + + # Handle variable exports: export const foo = ... + if child.type == "lexical_declaration": + for decl in child.children: + if decl.type == "variable_declarator": + name_node = decl.child_by_field_name("name") + if name_node and name_node.type == "identifier": + name = self.get_node_text(name_node, source_bytes) + exported_names.append((name, None)) + + # Skip if no exports found + if not exported_names and not default_export: + return None + + return ExportInfo( + exported_names=exported_names, + default_export=default_export, + is_reexport=is_reexport, + reexport_source=reexport_source, + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + wrapped_default_args=wrapped_default_args, + ) + + def _extract_call_expression_identifiers(self, node: Node, source_bytes: bytes) -> list[str]: + """Extract identifier names from arguments of a call expression. + + For patterns like curry(traverseEntity) or compose(fn1, fn2), this extracts + the function names passed as arguments: ["traverseEntity"] or ["fn1", "fn2"]. + + Args: + node: A call_expression node. + source_bytes: The source code as bytes. + + Returns: + List of identifier names found in the call arguments. + + """ + identifiers: list[str] = [] + + # Get the arguments node + args_node = node.child_by_field_name("arguments") + if args_node: + for child in args_node.children: + if child.type == "identifier": + identifiers.append(self.get_node_text(child, source_bytes)) + # Also handle nested call expressions: compose(curry(fn)) + elif child.type == "call_expression": + identifiers.extend(self._extract_call_expression_identifiers(child, source_bytes)) + + return identifiers + + def _extract_commonjs_export(self, node: Node, source_bytes: bytes) -> ExportInfo | None: + """Extract export information from CommonJS module.exports or exports.* patterns. + + Handles patterns like: + - module.exports = function() {} -> default export + - module.exports = { foo, bar } -> named exports + - module.exports.foo = function() {} -> named export 'foo' + - exports.foo = function() {} -> named export 'foo' + - module.exports = require('./other') -> re-export + """ + left_node = node.child_by_field_name("left") + right_node = node.child_by_field_name("right") + + if not left_node or not right_node: + return None + + # Check if this is a module.exports or exports.* pattern + if left_node.type != "member_expression": + return None + + left_text = self.get_node_text(left_node, source_bytes) + + exported_names: list[tuple[str, str | None]] = [] + default_export: str | None = None + is_reexport = False + reexport_source: str | None = None + + if left_text == "module.exports": + # module.exports = something + if right_node.type in {"function_expression", "arrow_function"}: + # module.exports = function foo() {} or module.exports = () => {} + name_node = right_node.child_by_field_name("name") + default_export = self.get_node_text(name_node, source_bytes) if name_node else "default" + elif right_node.type == "identifier": + # module.exports = someFunction + default_export = self.get_node_text(right_node, source_bytes) + elif right_node.type == "object": + # module.exports = { foo, bar, baz: qux } + for child in right_node.children: + if child.type == "shorthand_property_identifier": + # { foo } - exports function named foo + name = self.get_node_text(child, source_bytes) + exported_names.append((name, None)) + elif child.type == "pair": + # { baz: qux } - exports qux as baz + key_node = child.child_by_field_name("key") + value_node = child.child_by_field_name("value") + if key_node and value_node: + export_name = self.get_node_text(key_node, source_bytes) + local_name = self.get_node_text(value_node, source_bytes) + # In CommonJS { baz: qux }, baz is the exported name, qux is local + exported_names.append((local_name, export_name)) + elif right_node.type == "call_expression": + # module.exports = require('./other') - re-export + func_node = right_node.child_by_field_name("function") + if func_node and self.get_node_text(func_node, source_bytes) == "require": + is_reexport = True + args_node = right_node.child_by_field_name("arguments") + if args_node: + for arg in args_node.children: + if arg.type == "string": + reexport_source = self.get_node_text(arg, source_bytes).strip("'\"") + break + default_export = "default" + else: + # module.exports = something else (class, etc.) + default_export = "default" + + elif left_text.startswith("module.exports."): + # module.exports.foo = something + prop_name = left_text.split(".", 2)[2] # Get 'foo' from 'module.exports.foo' + exported_names.append((prop_name, None)) + + elif left_text.startswith("exports."): + # exports.foo = something + prop_name = left_text.split(".", 1)[1] # Get 'foo' from 'exports.foo' + exported_names.append((prop_name, None)) + + else: + # Not a CommonJS export pattern + return None + + # Skip if no exports found + if not exported_names and not default_export: + return None + + return ExportInfo( + exported_names=exported_names, + default_export=default_export, + is_reexport=is_reexport, + reexport_source=reexport_source, + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + ) + + def is_function_exported( + self, source: str, function_name: str, class_name: str | None = None + ) -> tuple[bool, str | None]: + """Check if a function is exported and get its export name. + + For class methods, also checks if the containing class is exported. + Also handles wrapped exports like: export default curry(traverseEntity) + + Args: + source: The source code to analyze. + function_name: The name of the function to check. + class_name: For class methods, the name of the containing class. + + Returns: + Tuple of (is_exported, export_name). export_name may differ from + function_name if exported with an alias. For class methods, + returns the class export name. + + """ + exports = self.find_exports(source) + + # First, check if the function itself is directly exported + for export in exports: + # Check default export + if export.default_export == function_name: + return (True, "default") + + # Check named exports + for name, alias in export.exported_names: + if name == function_name: + return (True, alias if alias else name) + + # Check wrapped default exports: export default curry(traverseEntity) + # The function is exported via wrapper, so it's accessible as "default" + if export.wrapped_default_args and function_name in export.wrapped_default_args: + return (True, "default") + + # For class methods, check if the containing class is exported + if class_name: + for export in exports: + # Check if class is default export + if export.default_export == class_name: + return (True, class_name) + + # Check if class is in named exports + for name, alias in export.exported_names: + if name == class_name: + return (True, alias if alias else name) + + return (False, None) + + def find_function_calls(self, source: str, within_function: FunctionNode) -> list[str]: + """Find all function calls within a specific function's body. + + Args: + source: The full source code. + within_function: The function to search within. + + Returns: + List of function names that are called. + + """ + calls: list[str] = [] + source_bytes = source.encode("utf8") + + # Get the body of the function + body_node = within_function.node.child_by_field_name("body") + if body_node is None: + # For arrow functions, the body might be the last child + for child in within_function.node.children: + if child.type in ("statement_block", "expression_statement") or ( + child.type not in ("identifier", "formal_parameters", "async", "=>") + ): + body_node = child + break + + if body_node: + self._walk_tree_for_calls(body_node, source_bytes, calls) + + return list(set(calls)) # Remove duplicates + + def _walk_tree_for_calls(self, node: Node, source_bytes: bytes, calls: list[str]) -> None: + """Recursively find function calls in a subtree.""" + if node.type == "call_expression": + func_node = node.child_by_field_name("function") + if func_node: + if func_node.type == "identifier": + calls.append(self.get_node_text(func_node, source_bytes)) + elif func_node.type == "member_expression": + # For method calls like obj.method(), get the method name + prop_node = func_node.child_by_field_name("property") + if prop_node: + calls.append(self.get_node_text(prop_node, source_bytes)) + + for child in node.children: + self._walk_tree_for_calls(child, source_bytes, calls) + + def find_module_level_declarations(self, source: str) -> list[ModuleLevelDeclaration]: + """Find all module-level variable/constant declarations. + + This finds global variables, constants, classes, enums, type aliases, + and interfaces defined at the top level of the module (not inside functions). + + Args: + source: The source code to analyze. + + Returns: + List of ModuleLevelDeclaration objects. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + declarations: list[ModuleLevelDeclaration] = [] + + # Only look at direct children of the program/module node (top-level) + for child in tree.root_node.children: + self._extract_module_level_declaration(child, source_bytes, declarations) + + return declarations + + def _extract_module_level_declaration( + self, node: Node, source_bytes: bytes, declarations: list[ModuleLevelDeclaration] + ) -> None: + """Extract module-level declarations from a node.""" + is_exported = False + + # Handle export statements - unwrap to get the actual declaration + if node.type == "export_statement": + is_exported = True + # Find the actual declaration inside the export + for child in node.children: + if child.type in ("lexical_declaration", "variable_declaration"): + self._extract_declaration(child, source_bytes, declarations, is_exported, node) + return + if child.type == "class_declaration": + name_node = child.child_by_field_name("name") + if name_node: + declarations.append( + ModuleLevelDeclaration( + name=self.get_node_text(name_node, source_bytes), + declaration_type="class", + source_code=self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + return + if child.type in ("type_alias_declaration", "interface_declaration", "enum_declaration"): + name_node = child.child_by_field_name("name") + if name_node: + decl_type = child.type.replace("_declaration", "").replace("_alias", "") + declarations.append( + ModuleLevelDeclaration( + name=self.get_node_text(name_node, source_bytes), + declaration_type=decl_type, + source_code=self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + return + return + + # Handle non-exported declarations + if node.type in ( + "lexical_declaration", # const/let + "variable_declaration", # var + ): + self._extract_declaration(node, source_bytes, declarations, is_exported, node) + elif node.type == "class_declaration": + name_node = node.child_by_field_name("name") + if name_node: + declarations.append( + ModuleLevelDeclaration( + name=self.get_node_text(name_node, source_bytes), + declaration_type="class", + source_code=self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + elif node.type in ("type_alias_declaration", "interface_declaration", "enum_declaration"): + name_node = node.child_by_field_name("name") + if name_node: + decl_type = node.type.replace("_declaration", "").replace("_alias", "") + declarations.append( + ModuleLevelDeclaration( + name=self.get_node_text(name_node, source_bytes), + declaration_type=decl_type, + source_code=self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + def _extract_declaration( + self, + node: Node, + source_bytes: bytes, + declarations: list[ModuleLevelDeclaration], + is_exported: bool, + source_node: Node, + ) -> None: + """Extract variable declarations (const/let/var).""" + # Determine declaration type (const, let, var) + decl_type = "var" + for child in node.children: + if child.type in ("const", "let", "var"): + decl_type = child.type + break + + # Find variable declarators + for child in node.children: + if child.type == "variable_declarator": + name_node = child.child_by_field_name("name") + if name_node: + # Handle destructuring patterns + if name_node.type == "identifier": + declarations.append( + ModuleLevelDeclaration( + name=self.get_node_text(name_node, source_bytes), + declaration_type=decl_type, + source_code=self.get_node_text(source_node, source_bytes), + start_line=source_node.start_point[0] + 1, + end_line=source_node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + elif name_node.type in ("object_pattern", "array_pattern"): + # For destructuring, extract all bound identifiers + identifiers = self._extract_pattern_identifiers(name_node, source_bytes) + for ident in identifiers: + declarations.append( + ModuleLevelDeclaration( + name=ident, + declaration_type=decl_type, + source_code=self.get_node_text(source_node, source_bytes), + start_line=source_node.start_point[0] + 1, + end_line=source_node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + def _extract_pattern_identifiers(self, pattern_node: Node, source_bytes: bytes) -> list[str]: + """Extract all identifier names from a destructuring pattern.""" + identifiers: list[str] = [] + + def walk(n: Node) -> None: + if n.type in {"identifier", "shorthand_property_identifier_pattern"}: + identifiers.append(self.get_node_text(n, source_bytes)) + for child in n.children: + walk(child) + + walk(pattern_node) + return identifiers + + def find_referenced_identifiers(self, source: str) -> set[str]: + """Find all identifiers referenced in the source code. + + This finds all identifier references, excluding: + - Declaration names (left side of assignments) + - Property names in object literals + - Function/class names at definition site + + Args: + source: The source code to analyze. + + Returns: + Set of referenced identifier names. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + references: set[str] = set() + + self._walk_tree_for_references(tree.root_node, source_bytes, references) + + return references + + def _walk_tree_for_references(self, node: Node, source_bytes: bytes, references: set[str]) -> None: + """Walk tree to collect identifier references.""" + if node.type == "identifier": + # Check if this identifier is a reference (not a declaration) + parent = node.parent + if parent is None: + return + + # Skip function/class/method names at definition + if parent.type in ("function_declaration", "class_declaration", "method_definition", "function_expression"): + if parent.child_by_field_name("name") == node: + # Don't recurse into parent's children - the parent will be visited separately + return + + # Skip variable declarator names (left side of declaration) + if parent.type == "variable_declarator" and parent.child_by_field_name("name") == node: + # Don't recurse - the value will be visited when we visit the declarator + return + + # Skip property names in object literals (keys) + if parent.type == "pair" and parent.child_by_field_name("key") == node: + # Don't recurse - the value will be visited when we visit the pair + return + + # Skip property access property names (obj.property - skip 'property') + if parent.type == "member_expression" and parent.child_by_field_name("property") == node: + # Don't recurse - the object will be visited when we visit the member_expression + return + + # Skip import specifier names + if parent.type in ("import_specifier", "import_clause", "namespace_import"): + return + + # Skip export specifier names + if parent.type == "export_specifier": + return + + # Skip parameter names in function definitions (but NOT default values) + if parent.type == "formal_parameters": + return + if parent.type == "required_parameter": + # Only skip if this is the parameter name (pattern field), not the default value + if parent.child_by_field_name("pattern") == node: + return + # If it's the value field (default value), it's a reference - don't skip + + # This is a reference + references.add(self.get_node_text(node, source_bytes)) + return + + # Recurse into children + for child in node.children: + self._walk_tree_for_references(child, source_bytes, references) + + def has_return_statement(self, function_node: FunctionNode, source: str) -> bool: + """Check if a function has a return statement. + + Args: + function_node: The function to check. + source: The source code. + + Returns: + True if the function has a return statement. + + """ + source_bytes = source.encode("utf8") + + # Generator functions always implicitly return a Generator/Iterator + if function_node.is_generator: + return True + + # For arrow functions with expression body, there's an implicit return + if function_node.is_arrow: + body_node = function_node.node.child_by_field_name("body") + if body_node and body_node.type != "statement_block": + # Expression body (implicit return) + return True + + return self._node_has_return(function_node.node) + + def _node_has_return(self, node: Node) -> bool: + """Recursively check if a node contains a return statement.""" + if node.type == "return_statement": + return True + + # Don't recurse into nested function definitions + if node.type in ("function_declaration", "function_expression", "arrow_function", "method_definition"): + # Only check the current function, not nested ones + body_node = node.child_by_field_name("body") + if body_node: + for child in body_node.children: + if self._node_has_return(child): + return True + return False + + return any(self._node_has_return(child) for child in node.children) + + def extract_type_annotations(self, source: str, function_name: str, function_line: int) -> set[str]: + """Extract type annotation names from a function's parameters and return type. + + Finds the function by name and line number, then extracts all user-defined type names + from its type annotations (parameters and return type). + + Args: + source: The source code to analyze. + function_name: Name of the function to find. + function_line: Start line of the function (1-indexed). + + Returns: + Set of type names found in the function's annotations. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + type_names: set[str] = set() + + # Find the function node + func_node = self._find_function_node(tree.root_node, source_bytes, function_name, function_line) + if not func_node: + return type_names + + # Extract type annotations from parameters + params_node = func_node.child_by_field_name("parameters") + if params_node: + self._extract_type_names_from_node(params_node, source_bytes, type_names) + + # Extract return type annotation + return_type_node = func_node.child_by_field_name("return_type") + if return_type_node: + self._extract_type_names_from_node(return_type_node, source_bytes, type_names) + + return type_names + + def extract_class_field_types(self, source: str, class_name: str) -> set[str]: + """Extract type annotation names from class field declarations. + + Args: + source: The source code to analyze. + class_name: Name of the class to analyze. + + Returns: + Set of type names found in class field annotations. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + type_names: set[str] = set() + + # Find the class node + class_node = self._find_class_node(tree.root_node, source_bytes, class_name) + if not class_node: + return type_names + + # Find class body and extract field type annotations + body_node = class_node.child_by_field_name("body") + if body_node: + for child in body_node.children: + # Handle public_field_definition (JS/TS class fields) + if child.type in ("public_field_definition", "field_definition"): + type_annotation = child.child_by_field_name("type") + if type_annotation: + self._extract_type_names_from_node(type_annotation, source_bytes, type_names) + + return type_names + + def _find_function_node( + self, node: Node, source_bytes: bytes, function_name: str, function_line: int + ) -> Node | None: + """Find a function/method node by name and line number.""" + if node.type in ( + "function_declaration", + "method_definition", + "function_expression", + "generator_function_declaration", + ): + name_node = node.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + # Line is 1-indexed, tree-sitter is 0-indexed + if name == function_name and (node.start_point[0] + 1) == function_line: + return node + + # Check arrow functions assigned to variables + if node.type == "lexical_declaration": + for child in node.children: + if child.type == "variable_declarator": + name_node = child.child_by_field_name("name") + value_node = child.child_by_field_name("value") + if name_node and value_node and value_node.type == "arrow_function": + name = self.get_node_text(name_node, source_bytes) + if name == function_name and (node.start_point[0] + 1) == function_line: + return value_node + + # Recurse into children + for child in node.children: + result = self._find_function_node(child, source_bytes, function_name, function_line) + if result: + return result + + return None + + def _find_class_node(self, node: Node, source_bytes: bytes, class_name: str) -> Node | None: + """Find a class node by name.""" + if node.type in ("class_declaration", "class"): + name_node = node.child_by_field_name("name") + if name_node: + name = self.get_node_text(name_node, source_bytes) + if name == class_name: + return node + + for child in node.children: + result = self._find_class_node(child, source_bytes, class_name) + if result: + return result + + return None + + def _extract_type_names_from_node(self, node: Node, source_bytes: bytes, type_names: set[str]) -> None: + """Recursively extract type names from a type annotation node. + + Handles various TypeScript type annotation patterns: + - Simple types: number, string, Point + - Generic types: Array, Promise + - Union types: A | B + - Intersection types: A & B + - Array types: T[] + - Tuple types: [A, B] + - Object/mapped types: { key: Type } + + Args: + node: Tree-sitter node to analyze. + source_bytes: Source code as bytes. + type_names: Set to add found type names to. + + """ + # Handle type identifiers (the actual type name references) + if node.type == "type_identifier": + type_name = self.get_node_text(node, source_bytes) + # Skip primitive types + if type_name not in ( + "number", + "string", + "boolean", + "void", + "null", + "undefined", + "any", + "never", + "unknown", + "object", + "symbol", + "bigint", + ): + type_names.add(type_name) + return + + # Handle regular identifiers in type position (can happen in some contexts) + if node.type == "identifier" and node.parent and node.parent.type in ("type_annotation", "generic_type"): + type_name = self.get_node_text(node, source_bytes) + if type_name not in ( + "number", + "string", + "boolean", + "void", + "null", + "undefined", + "any", + "never", + "unknown", + "object", + "symbol", + "bigint", + ): + type_names.add(type_name) + return + + # Handle nested_type_identifier (e.g., Namespace.Type) + if node.type == "nested_type_identifier": + # Get the full qualified name + type_name = self.get_node_text(node, source_bytes) + # Add both the full name and the first part (namespace) + type_names.add(type_name) + # Also extract the module/namespace part + module_node = node.child_by_field_name("module") + if module_node: + type_names.add(self.get_node_text(module_node, source_bytes)) + return + + # Recurse into all children for compound types + for child in node.children: + self._extract_type_names_from_node(child, source_bytes, type_names) + + def find_type_definitions(self, source: str) -> list[TypeDefinition]: + """Find all type definitions (interface, type, class, enum) in source code. + + Args: + source: The source code to analyze. + + Returns: + List of TypeDefinition objects. + + """ + source_bytes = source.encode("utf8") + tree = self.parse(source_bytes) + definitions: list[TypeDefinition] = [] + + # Walk through top-level nodes + for child in tree.root_node.children: + self._extract_type_definition(child, source_bytes, definitions) + + return definitions + + def _extract_type_definition( + self, node: Node, source_bytes: bytes, definitions: list[TypeDefinition], is_exported: bool = False + ) -> None: + """Extract type definitions from a node.""" + # Handle export statements - unwrap to get the actual definition + if node.type == "export_statement": + for child in node.children: + if child.type in ( + "interface_declaration", + "type_alias_declaration", + "class_declaration", + "enum_declaration", + ): + self._extract_type_definition(child, source_bytes, definitions, is_exported=True) + return + + # Extract interface definitions + if node.type == "interface_declaration": + name_node = node.child_by_field_name("name") + if name_node: + # Look for preceding JSDoc comment + jsdoc = "" + prev_sibling = node.prev_named_sibling + if prev_sibling and prev_sibling.type == "comment": + comment_text = self.get_node_text(prev_sibling, source_bytes) + if comment_text.strip().startswith("/**"): + jsdoc = comment_text + "\n" + + definitions.append( + TypeDefinition( + name=self.get_node_text(name_node, source_bytes), + definition_type="interface", + source_code=jsdoc + self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + # Extract type alias definitions + elif node.type == "type_alias_declaration": + name_node = node.child_by_field_name("name") + if name_node: + # Look for preceding JSDoc comment + jsdoc = "" + prev_sibling = node.prev_named_sibling + if prev_sibling and prev_sibling.type == "comment": + comment_text = self.get_node_text(prev_sibling, source_bytes) + if comment_text.strip().startswith("/**"): + jsdoc = comment_text + "\n" + + definitions.append( + TypeDefinition( + name=self.get_node_text(name_node, source_bytes), + definition_type="type", + source_code=jsdoc + self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + # Extract enum definitions + elif node.type == "enum_declaration": + name_node = node.child_by_field_name("name") + if name_node: + # Look for preceding JSDoc comment + jsdoc = "" + prev_sibling = node.prev_named_sibling + if prev_sibling and prev_sibling.type == "comment": + comment_text = self.get_node_text(prev_sibling, source_bytes) + if comment_text.strip().startswith("/**"): + jsdoc = comment_text + "\n" + + definitions.append( + TypeDefinition( + name=self.get_node_text(name_node, source_bytes), + definition_type="enum", + source_code=jsdoc + self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + # Extract class definitions (as types) + elif node.type == "class_declaration": + name_node = node.child_by_field_name("name") + if name_node: + # Look for preceding JSDoc comment + jsdoc = "" + prev_sibling = node.prev_named_sibling + if prev_sibling and prev_sibling.type == "comment": + comment_text = self.get_node_text(prev_sibling, source_bytes) + if comment_text.strip().startswith("/**"): + jsdoc = comment_text + "\n" + + definitions.append( + TypeDefinition( + name=self.get_node_text(name_node, source_bytes), + definition_type="class", + source_code=jsdoc + self.get_node_text(node, source_bytes), + start_line=node.start_point[0] + 1, + end_line=node.end_point[0] + 1, + is_exported=is_exported, + ) + ) + + +def get_analyzer_for_file(file_path: Path) -> TreeSitterAnalyzer: + """Get the appropriate TreeSitterAnalyzer for a file based on its extension. + + Args: + file_path: Path to the file. + + Returns: + TreeSitterAnalyzer configured for the file's language. + + """ + suffix = file_path.suffix.lower() + + if suffix == ".ts": + return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT) + if suffix == ".tsx": + return TreeSitterAnalyzer(TreeSitterLanguage.TSX) + # Default to JavaScript for .js, .jsx, .mjs, .cjs + return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT) + + +# Author: Saurabh Misra +def extract_calling_function_source(source_code: str, function_name: str, ref_line: int) -> str | None: + """Extract the source code of a calling function in JavaScript/TypeScript. + + Args: + source_code: Full source code of the file. + function_name: Name of the function to extract. + ref_line: Line number where the reference is (helps identify the right function). + + Returns: + Source code of the function, or None if not found. + + """ + try: + from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage + + # Try TypeScript first, fall back to JavaScript + for lang in [TreeSitterLanguage.TYPESCRIPT, TreeSitterLanguage.TSX, TreeSitterLanguage.JAVASCRIPT]: + try: + analyzer = TreeSitterAnalyzer(lang) + functions = analyzer.find_functions(source_code, include_methods=True) + + for func in functions: + if func.name == function_name: + # Check if the reference line is within this function + if func.start_line <= ref_line <= func.end_line: + return func.source_text + break + except Exception: + continue + + return None + except Exception: + return None diff --git a/codeflash/languages/javascript/vitest_runner.py b/codeflash/languages/javascript/vitest_runner.py index f622d7384..d169752bc 100644 --- a/codeflash/languages/javascript/vitest_runner.py +++ b/codeflash/languages/javascript/vitest_runner.py @@ -23,8 +23,13 @@ def _find_vitest_project_root(file_path: Path) -> Path | None: """Find the Vitest project root by looking for vitest/vite config or package.json. - Traverses up from the given file path to find the nearest directory - containing vitest.config.js/ts, vite.config.js/ts, or package.json. + Traverses up from the given file path to find the directory containing + vitest.config.js/ts or vite.config.js/ts. Falls back to package.json only + if no vitest/vite config is found in any parent directory. + + In monorepos, package.json may exist at multiple levels (e.g., packages/lib/package.json), + but the vitest config with setupFiles is typically at the monorepo root. + We need to prioritize finding the actual vitest config to ensure paths resolve correctly. Args: file_path: A file path within the Vitest project. @@ -34,8 +39,10 @@ def _find_vitest_project_root(file_path: Path) -> Path | None: """ current = file_path.parent if file_path.is_file() else file_path + package_json_dir = None # Track first package.json found (fallback) + while current != current.parent: # Stop at filesystem root - # Check for Vitest-specific config files first + # Check for Vitest-specific config files first - these should take priority if ( (current / "vitest.config.js").exists() or (current / "vitest.config.ts").exists() @@ -45,27 +52,40 @@ def _find_vitest_project_root(file_path: Path) -> Path | None: or (current / "vite.config.ts").exists() or (current / "vite.config.mjs").exists() or (current / "vite.config.mts").exists() - or (current / "package.json").exists() ): return current + # Remember first package.json as fallback, but keep looking for vitest config + if package_json_dir is None and (current / "package.json").exists(): + package_json_dir = current current = current.parent - return None + + # No vitest config found, fall back to package.json directory if found + return package_json_dir def _is_vitest_coverage_available(project_root: Path) -> bool: """Check if Vitest coverage package is available. + In monorepos, dependencies may be hoisted to the root node_modules. + This function searches up the directory tree for the coverage package. + Args: - project_root: The project root directory. + project_root: The project root directory (may be a package in a monorepo). Returns: True if @vitest/coverage-v8 or @vitest/coverage-istanbul is installed. """ - node_modules = project_root / "node_modules" - return (node_modules / "@vitest" / "coverage-v8").exists() or ( - node_modules / "@vitest" / "coverage-istanbul" - ).exists() + current = project_root + while current != current.parent: # Stop at filesystem root + node_modules = current / "node_modules" + if node_modules.exists(): + if (node_modules / "@vitest" / "coverage-v8").exists() or ( + node_modules / "@vitest" / "coverage-istanbul" + ).exists(): + return True + current = current.parent + return False def _ensure_runtime_files(project_root: Path) -> None: @@ -97,8 +117,146 @@ def _ensure_runtime_files(project_root: Path) -> None: logger.error(f"Could not install codeflash. Please install it manually: {' '.join(install_cmd)}") +def _find_monorepo_root(start_path: Path) -> Path | None: + """Find the monorepo root by looking for workspace markers. + + Args: + start_path: A path within the monorepo. + + Returns: + The monorepo root directory, or None if not found. + + """ + monorepo_markers = ["pnpm-workspace.yaml", "yarn.lock", "lerna.json", "package-lock.json"] + current = start_path if start_path.is_dir() else start_path.parent + + while current != current.parent: + # Check for monorepo markers + if any((current / marker).exists() for marker in monorepo_markers): + # Verify it has node_modules or package.json (it's a real root) + if (current / "node_modules").exists() or (current / "package.json").exists(): + return current + current = current.parent + + return None + + +def _is_vitest_workspace(project_root: Path) -> bool: + """Check if the project uses vitest workspace configuration. + + Vitest workspaces have a special structure where the root config + points to package-level configs. We shouldn't override these. + + Args: + project_root: The project root directory. + + Returns: + True if the project appears to use vitest workspace. + + """ + vitest_config = project_root / "vitest.config.ts" + if not vitest_config.exists(): + vitest_config = project_root / "vitest.config.js" + if not vitest_config.exists(): + return False + + try: + content = vitest_config.read_text() + # Check for workspace indicators + return "workspace" in content.lower() or "defineWorkspace" in content + except Exception: + return False + + +def _ensure_codeflash_vitest_config(project_root: Path) -> Path | None: + """Create or find a Codeflash-compatible Vitest config. + + Vitest configs often have restrictive include patterns like 'test/**/*.test.ts' + which filter out our generated test files. This function creates a config + that overrides the include pattern to accept all test files. + + Note: For workspace projects, we skip creating a custom config as it would + conflict with the workspace setup. In those cases, tests should be placed + in the correct package's test directory. + + Args: + project_root: The project root directory. + + Returns: + Path to the Codeflash Vitest config, or None if creation failed/not needed. + + """ + # Check for workspace configuration - don't override these + monorepo_root = _find_monorepo_root(project_root) + if monorepo_root and _is_vitest_workspace(monorepo_root): + logger.debug("Detected vitest workspace configuration - skipping custom config") + return None + + codeflash_config_path = project_root / "codeflash.vitest.config.mjs" + + # If already exists, use it + if codeflash_config_path.exists(): + logger.debug(f"Using existing Codeflash Vitest config: {codeflash_config_path}") + return codeflash_config_path + + # Find the original vitest config to extend + original_config = None + for config_name in ["vitest.config.ts", "vitest.config.js", "vitest.config.mts", "vitest.config.mjs"]: + config_path = project_root / config_name + if config_path.exists(): + original_config = config_name + break + + # Also check for vite config with vitest settings + if not original_config: + for config_name in ["vite.config.ts", "vite.config.js", "vite.config.mts", "vite.config.mjs"]: + config_path = project_root / config_name + if config_path.exists(): + original_config = config_name + break + + # Create a config that extends the original and overrides include pattern + if original_config: + config_content = f"""// Auto-generated by Codeflash for test file pattern compatibility +import {{ mergeConfig }} from 'vitest/config'; +import originalConfig from './{original_config}'; + +export default mergeConfig(originalConfig, {{ + test: {{ + // Override include pattern to match all test files including generated ones + include: ['**/*.test.ts', '**/*.test.js', '**/*.test.tsx', '**/*.test.jsx'], + }}, +}}); +""" + else: + # No original config found, create a minimal one + config_content = """// Auto-generated by Codeflash for test file pattern compatibility +import { defineConfig } from 'vitest/config'; + +export default defineConfig({ + test: { + // Include all test files including generated ones + include: ['**/*.test.ts', '**/*.test.js', '**/*.test.tsx', '**/*.test.jsx'], + // Exclude common non-test directories + exclude: ['**/node_modules/**', '**/dist/**'], + }, +}); +""" + + try: + codeflash_config_path.write_text(config_content) + logger.debug(f"Created Codeflash Vitest config: {codeflash_config_path}") + return codeflash_config_path + except Exception as e: + logger.warning(f"Failed to create Codeflash Vitest config: {e}") + return None + + def _build_vitest_behavioral_command( - test_files: list[Path], timeout: int | None = None, output_file: Path | None = None + test_files: list[Path], + timeout: int | None = None, + output_file: Path | None = None, + project_root: Path | None = None, ) -> list[str]: """Build Vitest command for behavioral tests. @@ -106,6 +264,7 @@ def _build_vitest_behavioral_command( test_files: List of test files to run. timeout: Optional timeout in seconds. output_file: Optional path for JUnit XML output. + project_root: Project root directory for --root flag. Returns: Command list for subprocess execution. @@ -120,6 +279,14 @@ def _build_vitest_behavioral_command( "--no-file-parallelism", # Serial execution for deterministic timing ] + # For monorepos with restrictive vitest configs (e.g., include: test/**/*.test.ts), + # we need to create a custom config that allows all test patterns. + # This is done by creating a codeflash.vitest.config.mjs file. + if project_root: + codeflash_vitest_config = _ensure_codeflash_vitest_config(project_root) + if codeflash_vitest_config: + cmd.append(f"--config={codeflash_vitest_config}") + if output_file: # Use dot notation for junit reporter output file when multiple reporters are used # Format: --outputFile.junit=/path/to/file.xml @@ -135,7 +302,10 @@ def _build_vitest_behavioral_command( def _build_vitest_benchmarking_command( - test_files: list[Path], timeout: int | None = None, output_file: Path | None = None + test_files: list[Path], + timeout: int | None = None, + output_file: Path | None = None, + project_root: Path | None = None, ) -> list[str]: """Build Vitest command for benchmarking tests. @@ -143,6 +313,7 @@ def _build_vitest_benchmarking_command( test_files: List of test files to run. timeout: Optional timeout in seconds. output_file: Optional path for JUnit XML output. + project_root: Project root directory for --root flag. Returns: Command list for subprocess execution. @@ -157,6 +328,12 @@ def _build_vitest_benchmarking_command( "--no-file-parallelism", # Serial execution for consistent benchmarking ] + # Use codeflash vitest config to override restrictive include patterns + if project_root: + codeflash_vitest_config = _ensure_codeflash_vitest_config(project_root) + if codeflash_vitest_config: + cmd.append(f"--config={codeflash_vitest_config}") + if output_file: # Use dot notation for junit reporter output file when multiple reporters are used cmd.append(f"--outputFile.junit={output_file}") @@ -220,11 +397,20 @@ def run_vitest_behavioral_tests( logger.debug("Vitest coverage package not installed, running without coverage") # Build Vitest command - vitest_cmd = _build_vitest_behavioral_command(test_files=test_files, timeout=timeout, output_file=result_file_path) + vitest_cmd = _build_vitest_behavioral_command( + test_files=test_files, timeout=timeout, output_file=result_file_path, project_root=effective_cwd + ) # Add coverage flags only if coverage is available if coverage_available: + # Don't pre-create the coverage directory - vitest should create it + # Pre-creating an empty directory may cause vitest to delete it + logger.debug(f"Coverage will be written to: {coverage_dir}") + vitest_cmd.extend(["--coverage", "--coverage.reporter=json", f"--coverage.reportsDirectory={coverage_dir}"]) + # Note: Removed --coverage.enabled=true (redundant) and --coverage.all false + # The version mismatch between vitest and @vitest/coverage-v8 can cause + # issues with coverage flag parsing. Let vitest use default settings. # Set up environment vitest_env = test_env.copy() @@ -251,6 +437,7 @@ def run_vitest_behavioral_tests( cwd=effective_cwd, env=vitest_env, timeout=subprocess_timeout, check=False, text=True, capture_output=True ) result = subprocess.run(vitest_cmd, **run_args) # noqa: PLW1510 + # Combine stderr into stdout for timing markers if result.stderr and not result.stdout: result = subprocess.CompletedProcess( @@ -296,6 +483,26 @@ def run_vitest_behavioral_tests( f"Vitest stdout: {result.stdout[:1000] if result.stdout else '(empty)'}" ) + # Check if coverage file was created + if coverage_available and coverage_json_path: + if coverage_json_path.exists(): + cov_size = coverage_json_path.stat().st_size + logger.debug(f"Vitest coverage JSON created: {coverage_json_path} ({cov_size} bytes)") + else: + # Check if the parent directory exists and list its contents + cov_parent = coverage_json_path.parent + if cov_parent.exists(): + contents = list(cov_parent.iterdir()) + logger.warning( + f"Vitest coverage JSON not created at {coverage_json_path}. " + f"Directory exists with contents: {[f.name for f in contents]}" + ) + else: + logger.warning( + f"Vitest coverage JSON not created at {coverage_json_path}. " + f"Coverage directory does not exist: {cov_parent}" + ) + return result_file_path, result, coverage_json_path, None @@ -313,6 +520,9 @@ def run_vitest_benchmarking_tests( ) -> tuple[Path, subprocess.CompletedProcess]: """Run Vitest benchmarking tests with external looping from Python. + NOTE: This function MUST use benchmarking_file_path (perf tests with capturePerf), + NOT instrumented_behavior_file_path (behavior tests with capture). + Uses external process-level looping to run tests multiple times and collect timing data. This matches the Python pytest approach where looping is controlled externally for simplicity. @@ -337,6 +547,26 @@ def run_vitest_benchmarking_tests( # Get performance test files test_files = [Path(file.benchmarking_file_path) for file in test_paths.test_files if file.benchmarking_file_path] + # Log test file selection + total_test_files = len(test_paths.test_files) + perf_test_files = len(test_files) + logger.debug( + f"Vitest benchmark test file selection: {perf_test_files}/{total_test_files} have benchmarking_file_path" + ) + if perf_test_files == 0: + logger.warning("No perf test files found! Cannot run benchmarking tests.") + for tf in test_paths.test_files: + logger.warning( + f"Test file: behavior={tf.instrumented_behavior_file_path}, perf={tf.benchmarking_file_path}" + ) + elif perf_test_files < total_test_files: + for tf in test_paths.test_files: + if not tf.benchmarking_file_path: + logger.warning(f"Missing benchmarking_file_path: behavior={tf.instrumented_behavior_file_path}") + else: + for tf in test_files[:3]: # Log first 3 perf test files + logger.debug(f"Using perf test file: {tf}") + # Use provided project_root, or detect it as fallback if project_root is None and test_files: project_root = _find_vitest_project_root(test_files[0]) @@ -349,7 +579,7 @@ def run_vitest_benchmarking_tests( # Build Vitest command for performance tests vitest_cmd = _build_vitest_benchmarking_command( - test_files=test_files, timeout=timeout, output_file=result_file_path + test_files=test_files, timeout=timeout, output_file=result_file_path, project_root=effective_cwd ) # Base environment setup @@ -367,14 +597,25 @@ def run_vitest_benchmarking_tests( vitest_env["CODEFLASH_PERF_STABILITY_CHECK"] = "true" if stability_check else "false" vitest_env["CODEFLASH_LOOP_INDEX"] = "1" + # Set test module for marker identification (use first test file as reference) + if test_files: + test_module_path = str( + test_files[0].relative_to(effective_cwd) + if test_files[0].is_relative_to(effective_cwd) + else test_files[0].name + ) + vitest_env["CODEFLASH_TEST_MODULE"] = test_module_path + logger.debug(f"[VITEST-BENCH] Set CODEFLASH_TEST_MODULE={test_module_path}") + # Total timeout for the entire benchmark run total_timeout = max(120, (target_duration_ms // 1000) + 60, timeout or 120) - logger.debug(f"Running Vitest benchmarking tests: {' '.join(vitest_cmd)}") + logger.debug(f"[VITEST-BENCH] Running Vitest benchmarking tests: {' '.join(vitest_cmd)}") logger.debug( - f"Vitest benchmarking config: min_loops={min_loops}, max_loops={max_loops}, " + f"[VITEST-BENCH] Config: min_loops={min_loops}, max_loops={max_loops}, " f"target_duration={target_duration_ms}ms, stability_check={stability_check}" ) + logger.debug(f"[VITEST-BENCH] Environment: CODEFLASH_PERF_LOOP_COUNT={vitest_env.get('CODEFLASH_PERF_LOOP_COUNT')}") total_start_time = time.time() @@ -399,7 +640,27 @@ def run_vitest_benchmarking_tests( result = subprocess.CompletedProcess(args=vitest_cmd, returncode=-1, stdout="", stderr="Vitest not found") wall_clock_seconds = time.time() - total_start_time - logger.debug(f"Vitest benchmarking completed in {wall_clock_seconds:.2f}s") + logger.debug(f"[VITEST-BENCH] Completed in {wall_clock_seconds:.2f}s, returncode={result.returncode}") + + # Debug: Check for END markers with duration (perf test format) + if result.stdout: + import re + + perf_end_pattern = re.compile(r"!######[^:]+:[^:]+:[^:]+:(\d+):[^:]+:(\d+)######!") + perf_matches = list(perf_end_pattern.finditer(result.stdout)) + if perf_matches: + loop_indices = [int(m.group(1)) for m in perf_matches] + logger.debug( + f"[VITEST-BENCH] Found {len(perf_matches)} perf END markers in stdout, " + f"loop_index range: {min(loop_indices)}-{max(loop_indices)}" + ) + else: + logger.debug(f"[VITEST-BENCH] No perf END markers found in stdout (len={len(result.stdout)})") + # Check if there are behavior END markers instead + behavior_end_pattern = re.compile(r"!######[^:]+:[^:]+:[^:]+:\d+:[^#]+######!") + behavior_matches = list(behavior_end_pattern.finditer(result.stdout)) + if behavior_matches: + logger.debug(f"[VITEST-BENCH] Found {len(behavior_matches)} behavior END markers instead (no duration)") return result_file_path, result @@ -460,6 +721,12 @@ def run_vitest_line_profile_tests( "--no-file-parallelism", # Serial execution for consistent line profiling ] + # Use codeflash vitest config to override restrictive include patterns + if effective_cwd: + codeflash_vitest_config = _ensure_codeflash_vitest_config(effective_cwd) + if codeflash_vitest_config: + vitest_cmd.append(f"--config={codeflash_vitest_config}") + # Use dot notation for junit reporter output file when multiple reporters are used vitest_cmd.append(f"--outputFile.junit={result_file_path}") diff --git a/codeflash/languages/python/__init__.py b/codeflash/languages/python/__init__.py index e599d1431..939d5941f 100644 --- a/codeflash/languages/python/__init__.py +++ b/codeflash/languages/python/__init__.py @@ -5,6 +5,7 @@ to the LanguageSupport protocol. """ +from codeflash.languages.python.reference_graph import ReferenceGraph from codeflash.languages.python.support import PythonSupport -__all__ = ["PythonSupport"] +__all__ = ["PythonSupport", "ReferenceGraph"] diff --git a/codeflash/languages/python/context/__init__.py b/codeflash/languages/python/context/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py new file mode 100644 index 000000000..d75e9fc8e --- /dev/null +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -0,0 +1,1508 @@ +from __future__ import annotations + +import ast +import hashlib +import os +from collections import defaultdict, deque +from itertools import chain +from pathlib import Path +from typing import TYPE_CHECKING + +import libcst as cst + +from codeflash.cli_cmds.console import logger +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.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, + get_section_names, + is_assignment_used, + recurse_sections, + remove_unused_definitions_by_function_names, +) +from codeflash.languages.python.static_analysis.code_extractor import ( + add_needed_imports_from_module, + find_preexisting_objects, +) +from codeflash.models.models import ( + CodeContextType, + CodeOptimizationContext, + CodeString, + CodeStringsMarkdown, + FunctionSource, +) +from codeflash.optimization.function_context import belongs_to_function_qualified + +if TYPE_CHECKING: + from jedi.api.classes import Name + + from codeflash.languages.base import DependencyResolver, 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" +TESTGEN_LIMIT_ERROR = "Testgen code context has exceeded token limit, cannot proceed" + + +def build_testgen_context( + helpers_of_fto_dict: dict[Path, set[FunctionSource]], + helpers_of_helpers_dict: dict[Path, set[FunctionSource]], + project_root_path: Path, + *, + remove_docstrings: bool = False, + include_enrichment: bool = True, + function_to_optimize: FunctionToOptimize | None = None, +) -> CodeStringsMarkdown: + testgen_context = extract_code_markdown_context_from_files( + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=remove_docstrings, + code_context_type=CodeContextType.TESTGEN, + ) + + if include_enrichment: + enrichment = enrich_testgen_context(testgen_context, project_root_path) + if enrichment.code_strings: + testgen_context = CodeStringsMarkdown(code_strings=testgen_context.code_strings + enrichment.code_strings) + + if function_to_optimize is not None: + result = _parse_and_collect_imports(testgen_context) + existing_classes = collect_existing_class_names(result[0]) if result else set() + constructor_stubs = extract_parameter_type_constructors( + function_to_optimize, project_root_path, existing_classes + ) + if constructor_stubs.code_strings: + testgen_context = CodeStringsMarkdown( + code_strings=testgen_context.code_strings + constructor_stubs.code_strings + ) + + return testgen_context + + +def get_code_optimization_context( + function_to_optimize: FunctionToOptimize, + project_root_path: Path, + optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT, + testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, + call_graph: DependencyResolver | None = None, +) -> CodeOptimizationContext: + # Route to language-specific implementation for non-Python languages + if not is_python(): + return get_code_optimization_context_for_language( + function_to_optimize, project_root_path, optim_token_limit, testgen_token_limit + ) + + # Get FunctionSource representation of helpers of FTO + fto_input = {function_to_optimize.file_path: {function_to_optimize.qualified_name}} + if call_graph is not None: + helpers_of_fto_dict, helpers_of_fto_list = call_graph.get_callees(fto_input) + else: + helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi(fto_input, project_root_path) + + # Add function to optimize into helpers of FTO dict, as they'll be processed together + fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path) + helpers_of_fto_dict[function_to_optimize.file_path].add(fto_as_function_source) + + # Format data to search for helpers of helpers using get_function_sources_from_jedi + helpers_of_fto_qualified_names_dict = { + file_path: {source.qualified_name for source in sources} for file_path, sources in helpers_of_fto_dict.items() + } + + # __init__ functions are automatically considered as helpers of FTO, so we add them to the dict (regardless of whether they exist) + # This helps us to search for helpers of __init__ functions of classes that contain helpers of FTO + for qualified_names in helpers_of_fto_qualified_names_dict.values(): + qualified_names.update({f"{qn.rsplit('.', 1)[0]}.__init__" for qn in qualified_names if "." in qn}) + + helpers_of_helpers_dict, helpers_of_helpers_list = get_function_sources_from_jedi( + helpers_of_fto_qualified_names_dict, project_root_path + ) + + # Extract code context for optimization + final_read_writable_code = extract_code_markdown_context_from_files( + helpers_of_fto_dict, + {}, + project_root_path, + remove_docstrings=False, + code_context_type=CodeContextType.READ_WRITABLE, + ) + + read_only_code_markdown = extract_code_markdown_context_from_files( + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=False, + code_context_type=CodeContextType.READ_ONLY, + ) + hashing_code_context = extract_code_markdown_context_from_files( + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=True, + code_context_type=CodeContextType.HASHING, + ) + + # Handle token limits + final_read_writable_tokens = encoded_tokens_len(final_read_writable_code.markdown) + if final_read_writable_tokens > optim_token_limit: + raise ValueError(READ_WRITABLE_LIMIT_ERROR) + + # Setup preexisting objects for code replacer + preexisting_objects = set( + chain( + *(find_preexisting_objects(codestring.code) for codestring in final_read_writable_code.code_strings), + *(find_preexisting_objects(codestring.code) for codestring in read_only_code_markdown.code_strings), + ) + ) + read_only_context_code = read_only_code_markdown.markdown + + # Progressive fallback for read-only context token limits + read_only_tokens = encoded_tokens_len(read_only_context_code) + if final_read_writable_tokens + read_only_tokens > optim_token_limit: + logger.debug("Code context has exceeded token limit, removing docstrings from read-only code") + read_only_code_no_docstrings = extract_code_markdown_context_from_files( + helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, remove_docstrings=True + ) + read_only_context_code = read_only_code_no_docstrings.markdown + if final_read_writable_tokens + encoded_tokens_len(read_only_context_code) > optim_token_limit: + logger.debug("Code context has exceeded token limit, removing read-only code") + read_only_context_code = "" + + # Progressive fallback for testgen context token limits + testgen_context = build_testgen_context( + helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, function_to_optimize=function_to_optimize + ) + + if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: + logger.debug("Testgen context exceeded token limit, removing docstrings") + testgen_context = build_testgen_context( + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=True, + function_to_optimize=function_to_optimize, + ) + + if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: + logger.debug("Testgen context still exceeded token limit, removing enrichment") + testgen_context = build_testgen_context( + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=True, + include_enrichment=False, + ) + + if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: + raise ValueError(TESTGEN_LIMIT_ERROR) + code_hash_context = hashing_code_context.markdown + code_hash = hashlib.sha256(code_hash_context.encode("utf-8")).hexdigest() + + all_helper_fqns = list({fs.fully_qualified_name for fs in helpers_of_fto_list + helpers_of_helpers_list}) + + return CodeOptimizationContext( + testgen_context=testgen_context, + read_writable_code=final_read_writable_code, + read_only_context_code=read_only_context_code, + hashing_code_context=code_hash_context, + hashing_code_context_hash=code_hash, + helper_functions=helpers_of_fto_list, + testgen_helper_fqns=all_helper_fqns, + preexisting_objects=preexisting_objects, + ) + + +def get_code_optimization_context_for_language( + function_to_optimize: FunctionToOptimize, + project_root_path: Path, + optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT, + testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, +) -> CodeOptimizationContext: + """Extract code optimization context for non-Python languages. + + Uses the language support abstraction to extract code context and converts + it to the CodeOptimizationContext format expected by the pipeline. + + This function supports multi-file context extraction, grouping helpers by file + and creating proper CodeStringsMarkdown with file paths for multi-file replacement. + + Args: + function_to_optimize: The function to extract context for. + project_root_path: Root of the project. + optim_token_limit: Token limit for optimization context. + testgen_token_limit: Token limit for testgen context. + + Returns: + CodeOptimizationContext with target code and dependencies. + + """ + from codeflash.languages import get_language_support + + # Get language support for this function + language = Language(function_to_optimize.language) + lang_support = get_language_support(language) + + # Extract code context using language support + code_context = lang_support.extract_code_context(function_to_optimize, project_root_path, project_root_path) + + # Build imports string if available + imports_code = "\n".join(code_context.imports) if code_context.imports else "" + + # Get relative path for target file + try: + target_relative_path = function_to_optimize.file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + target_relative_path = function_to_optimize.file_path + + # Group helpers by file path + helpers_by_file: dict[Path, list[HelperFunction]] = defaultdict(list) + helper_function_sources = [] + + for helper in code_context.helper_functions: + helpers_by_file[helper.file_path].append(helper) + + # Convert to FunctionSource for pipeline compatibility + helper_function_sources.append( + FunctionSource( + file_path=helper.file_path, + qualified_name=helper.qualified_name, + fully_qualified_name=helper.qualified_name, + only_function_name=helper.name, + source_code=helper.source_code, + ) + ) + + # Build read-writable code (target file + same-file helpers + global variables) + read_writable_code_strings = [] + + # Combine target code with same-file helpers + target_file_code = code_context.target_code + same_file_helpers = helpers_by_file.get(function_to_optimize.file_path, []) + if same_file_helpers: + helper_code = "\n\n".join(h.source_code for h in same_file_helpers) + target_file_code = target_file_code + "\n\n" + helper_code + + # Note: code_context.read_only_context contains type definitions and global variables + # These should be passed as read-only context to the AI, not prepended to the target code + # If prepended to target code, the AI treats them as code to optimize and includes them in output + + # Add imports to target file code + if imports_code: + target_file_code = imports_code + "\n\n" + target_file_code + + read_writable_code_strings.append( + CodeString(code=target_file_code, file_path=target_relative_path, language=function_to_optimize.language) + ) + + # Add helper files (cross-file helpers) + for file_path, file_helpers in helpers_by_file.items(): + if file_path == function_to_optimize.file_path: + continue # Already included in target file + + try: + helper_relative_path = file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + helper_relative_path = file_path + + # Combine all helpers from this file + combined_helper_code = "\n\n".join(h.source_code for h in file_helpers) + + read_writable_code_strings.append( + CodeString( + code=combined_helper_code, file_path=helper_relative_path, language=function_to_optimize.language + ) + ) + + read_writable_code = CodeStringsMarkdown( + code_strings=read_writable_code_strings, language=function_to_optimize.language + ) + + # Build testgen context (same as read_writable for non-Python) + testgen_context = CodeStringsMarkdown( + code_strings=read_writable_code_strings.copy(), language=function_to_optimize.language + ) + + # Check token limits + read_writable_tokens = encoded_tokens_len(read_writable_code.markdown) + if read_writable_tokens > optim_token_limit: + raise ValueError(READ_WRITABLE_LIMIT_ERROR) + + testgen_tokens = encoded_tokens_len(testgen_context.markdown) + if testgen_tokens > testgen_token_limit: + raise ValueError(TESTGEN_LIMIT_ERROR) + + # Generate code hash from all read-writable code + code_hash = hashlib.sha256(read_writable_code.flat.encode("utf-8")).hexdigest() + + return CodeOptimizationContext( + testgen_context=testgen_context, + read_writable_code=read_writable_code, + read_only_context_code=code_context.read_only_context, + hashing_code_context=read_writable_code.flat, + hashing_code_context_hash=code_hash, + helper_functions=helper_function_sources, + testgen_helper_fqns=[fs.fully_qualified_name for fs in helper_function_sources], + preexisting_objects=set(), + ) + + +def process_file_context( + file_path: Path, + primary_qualified_names: set[str], + secondary_qualified_names: set[str], + code_context_type: CodeContextType, + remove_docstrings: bool, + project_root_path: Path, + helper_functions: list[FunctionSource], +) -> CodeString | None: + try: + original_code = file_path.read_text("utf8") + except Exception as e: + logger.exception(f"Error while parsing {file_path}: {e}") + return None + + try: + all_names = primary_qualified_names | secondary_qualified_names + code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, all_names) + pruned_module = parse_code_and_prune_cst( + code_without_unused_defs, + code_context_type, + primary_qualified_names, + secondary_qualified_names, + remove_docstrings, + ) + except ValueError as e: + logger.debug(f"Error while getting read-only code: {e}") + return None + + if pruned_module.code.strip(): + if code_context_type == CodeContextType.HASHING: + code_context = ast.unparse(ast.parse(pruned_module.code)) + else: + code_context = add_needed_imports_from_module( + src_module_code=original_code, + dst_module_code=pruned_module, + src_path=file_path, + dst_path=file_path, + project_root=project_root_path, + helper_functions=helper_functions, + ) + try: + relative_path = file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + relative_path = file_path + return CodeString(code=code_context, file_path=relative_path) + return None + + +def extract_code_markdown_context_from_files( + helpers_of_fto: dict[Path, set[FunctionSource]], + helpers_of_helpers: dict[Path, set[FunctionSource]], + project_root_path: Path, + remove_docstrings: bool = False, + code_context_type: CodeContextType = CodeContextType.READ_ONLY, +) -> CodeStringsMarkdown: + """Extract code context from files containing target functions and their helpers, formatting them as markdown. + + This function processes two sets of files: + 1. Files containing the function to optimize (fto) and their first-degree helpers + 2. Files containing only helpers of helpers (with no overlap with the first set) + + For each file, it extracts relevant code based on the specified context type, adds necessary + imports, and combines them into a structured markdown format. + + Args: + ---- + helpers_of_fto: Dictionary mapping file paths to sets of Function Sources of function to optimize and its helpers + helpers_of_helpers: Dictionary mapping file paths to sets of Function Sources of helpers of helper functions + project_root_path: Root path of the project + remove_docstrings: Whether to remove docstrings from the extracted code + code_context_type: Type of code context to extract (READ_ONLY, READ_WRITABLE, or TESTGEN) + + Returns: + ------- + CodeStringsMarkdown containing the extracted code context with necessary imports, + formatted for inclusion in markdown + + """ + # Rearrange to remove overlaps, so we only access each file path once + helpers_of_helpers_no_overlap = defaultdict(set) + for file_path, function_sources in helpers_of_helpers.items(): + if file_path in helpers_of_fto: + # Remove duplicates within the same file path, in case a helper of helper is also a helper of fto + helpers_of_helpers[file_path] -= helpers_of_fto[file_path] + else: + helpers_of_helpers_no_overlap[file_path] = function_sources + code_context_markdown = CodeStringsMarkdown() + # Extract code from file paths that contain fto and first degree helpers. helpers of helpers may also be included if they are in the same files + for file_path, function_sources in helpers_of_fto.items(): + qualified_function_names = {func.qualified_name for func in function_sources} + helpers_of_helpers_qualified_names = {func.qualified_name for func in helpers_of_helpers.get(file_path, set())} + helper_functions = list(helpers_of_fto.get(file_path, set()) | helpers_of_helpers.get(file_path, set())) + + result = process_file_context( + file_path=file_path, + primary_qualified_names=qualified_function_names, + secondary_qualified_names=helpers_of_helpers_qualified_names, + code_context_type=code_context_type, + remove_docstrings=remove_docstrings, + project_root_path=project_root_path, + helper_functions=helper_functions, + ) + + if result is not None: + code_context_markdown.code_strings.append(result) + # Extract code from file paths containing helpers of helpers + for file_path, helper_function_sources in helpers_of_helpers_no_overlap.items(): + qualified_helper_function_names = {func.qualified_name for func in helper_function_sources} + helper_functions = list(helpers_of_helpers_no_overlap.get(file_path, set())) + + result = process_file_context( + file_path=file_path, + primary_qualified_names=set(), + secondary_qualified_names=qualified_helper_function_names, + code_context_type=code_context_type, + remove_docstrings=remove_docstrings, + project_root_path=project_root_path, + helper_functions=helper_functions, + ) + + if result is not None: + code_context_markdown.code_strings.append(result) + return code_context_markdown + + +def get_function_to_optimize_as_function_source( + function_to_optimize: FunctionToOptimize, project_root_path: Path +) -> FunctionSource: + import jedi + + # Use jedi to find function to optimize + script = jedi.Script(path=function_to_optimize.file_path, project=jedi.Project(path=project_root_path)) + + # Get all names in the file + names = script.get_names(all_scopes=True, definitions=True, references=False) + + # Find the name that matches our function + for name in names: + try: + if ( + name.type == "function" + and name.full_name + and name.name == function_to_optimize.function_name + and name.full_name.startswith(name.module_name) + and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name + ): + return FunctionSource( + file_path=function_to_optimize.file_path, + qualified_name=function_to_optimize.qualified_name, + fully_qualified_name=name.full_name, + only_function_name=name.name, + source_code=name.get_line_code(), + ) + except Exception as e: + logger.exception(f"Error while getting function source: {e}") + continue + raise ValueError( + f"Could not find function {function_to_optimize.function_name} in {function_to_optimize.file_path}" # noqa: EM102 + ) + + +def get_function_sources_from_jedi( + file_path_to_qualified_function_names: dict[Path, set[str]], project_root_path: Path +) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: + import jedi + + file_path_to_function_source = defaultdict(set) + function_source_list: list[FunctionSource] = [] + for file_path, qualified_function_names in file_path_to_qualified_function_names.items(): + script = jedi.Script(path=file_path, project=jedi.Project(path=project_root_path)) + file_refs = script.get_names(all_scopes=True, definitions=False, references=True) + + for qualified_function_name in qualified_function_names: + names = [ + ref + for ref in file_refs + if ref.full_name and belongs_to_function_qualified(ref, qualified_function_name) + ] + for name in names: + try: + definitions: list[Name] = name.goto(follow_imports=True, follow_builtin_imports=False) + except Exception: + logger.debug(f"Error while getting definitions for {qualified_function_name}") + definitions = [] + if definitions: + # TODO: there can be multiple definitions, see how to handle such cases + definition = definitions[0] + definition_path = definition.module_path + if definition_path is not None: + try: + rel = definition_path.resolve().relative_to(project_root_path.resolve()) + definition_path = project_root_path / rel + except ValueError: + pass + + # The definition is part of this project and not defined within the original function + is_valid_definition = ( + definition_path is not None + and not path_belongs_to_site_packages(definition_path) + and str(definition_path).startswith(str(project_root_path) + os.sep) + and definition.full_name + and not belongs_to_function_qualified(definition, qualified_function_name) + and definition.full_name.startswith(definition.module_name) + ) + if is_valid_definition and definition.type in ("function", "class", "statement"): + if definition.type == "function": + fqn = definition.full_name + func_name = definition.name + elif definition.type == "class": + fqn = f"{definition.full_name}.__init__" + func_name = "__init__" + else: + fqn = definition.full_name + func_name = definition.name + qualified_name = get_qualified_name(definition.module_name, fqn) + # Avoid nested functions or classes. Only class.function is allowed + if len(qualified_name.split(".")) <= 2: + function_source = FunctionSource( + file_path=definition_path, + qualified_name=qualified_name, + fully_qualified_name=fqn, + only_function_name=func_name, + source_code=definition.get_line_code(), + ) + file_path_to_function_source[definition_path].add(function_source) + function_source_list.append(function_source) + + return file_path_to_function_source, function_source_list + + +def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.Module, dict[str, str]] | None: + all_code = "\n".join(cs.code for cs in code_context.code_strings) + try: + tree = ast.parse(all_code) + except SyntaxError: + return None + imported_names: dict[str, str] = {} + + # Directly iterate over the module body and nested structures instead of ast.walk + # This avoids traversing every single node in the tree + def collect_imports(nodes: list[ast.stmt]) -> None: + for node in nodes: + if isinstance(node, ast.ImportFrom) and node.module: + for alias in node.names: + if alias.name != "*": + imported_name = alias.asname if alias.asname else alias.name + imported_names[imported_name] = node.module + # Recursively check nested structures (function defs, class defs, if statements, etc.) + elif isinstance( + node, + ( + ast.FunctionDef, + ast.AsyncFunctionDef, + ast.ClassDef, + ast.If, + ast.For, + ast.AsyncFor, + ast.While, + ast.With, + ast.AsyncWith, + ast.Try, + ast.ExceptHandler, + ), + ): + if hasattr(node, "body"): + collect_imports(node.body) + if hasattr(node, "orelse"): + collect_imports(node.orelse) + if hasattr(node, "finalbody"): + collect_imports(node.finalbody) + if hasattr(node, "handlers"): + for handler in node.handlers: + collect_imports(handler.body) + # Handle match/case statements (Python 3.10+) + elif hasattr(ast, "Match") and isinstance(node, ast.Match): + for case in node.cases: + collect_imports(case.body) + + collect_imports(tree.body) + return tree, imported_names + + +def collect_existing_class_names(tree: ast.Module) -> set[str]: + class_names = set() + stack = list(tree.body) + + while stack: + node = stack.pop() + if isinstance(node, ast.ClassDef): + class_names.add(node.name) + stack.extend(node.body) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + stack.extend(node.body) + elif isinstance(node, (ast.If, ast.For, ast.While, ast.With)): + stack.extend(node.body) + if hasattr(node, "orelse"): + stack.extend(node.orelse) + elif isinstance(node, ast.Try): + stack.extend(node.body) + stack.extend(node.orelse) + stack.extend(node.finalbody) + for handler in node.handlers: + stack.extend(handler.body) + + return class_names + + +BUILTIN_AND_TYPING_NAMES = frozenset( + { + "int", + "str", + "float", + "bool", + "bytes", + "bytearray", + "complex", + "list", + "dict", + "set", + "frozenset", + "tuple", + "type", + "object", + "None", + "NoneType", + "Ellipsis", + "NotImplemented", + "memoryview", + "range", + "slice", + "property", + "classmethod", + "staticmethod", + "super", + "Optional", + "Union", + "Any", + "List", + "Dict", + "Set", + "FrozenSet", + "Tuple", + "Type", + "Callable", + "Iterator", + "Generator", + "Coroutine", + "AsyncGenerator", + "AsyncIterator", + "Iterable", + "AsyncIterable", + "Sequence", + "MutableSequence", + "Mapping", + "MutableMapping", + "Collection", + "Awaitable", + "Literal", + "Final", + "ClassVar", + "TypeVar", + "TypeAlias", + "ParamSpec", + "Concatenate", + "Annotated", + "TypeGuard", + "Self", + "Unpack", + "TypeVarTuple", + "Never", + "NoReturn", + "SupportsInt", + "SupportsFloat", + "SupportsComplex", + "SupportsBytes", + "SupportsAbs", + "SupportsRound", + "IO", + "TextIO", + "BinaryIO", + "Pattern", + "Match", + } +) + + +def collect_type_names_from_annotation(node: ast.expr | None) -> set[str]: + if node is None: + return set() + if isinstance(node, ast.Name): + return {node.id} + if isinstance(node, ast.Subscript): + names = collect_type_names_from_annotation(node.value) + names |= collect_type_names_from_annotation(node.slice) + return names + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + return collect_type_names_from_annotation(node.left) | collect_type_names_from_annotation(node.right) + if isinstance(node, ast.Tuple): + names = set[str]() + for elt in node.elts: + names |= collect_type_names_from_annotation(elt) + return names + return set() + + +def extract_init_stub_from_class(class_name: str, module_source: str, module_tree: ast.Module) -> str | None: + class_node = None + # Use a deque-based BFS to find the first matching ClassDef (preserves ast.walk order) + q: deque[ast.AST] = deque([module_tree]) + while q: + candidate = q.popleft() + if isinstance(candidate, ast.ClassDef) and candidate.name == class_name: + class_node = candidate + break + q.extend(ast.iter_child_nodes(candidate)) + + if class_node is None: + return None + + lines = module_source.splitlines() + relevant_nodes: list[ast.FunctionDef | ast.AsyncFunctionDef] = [] + for item in class_node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)): + is_relevant = False + if item.name in ("__init__", "__post_init__"): + is_relevant = True + else: + # Check decorators explicitly to avoid generator overhead + for d in item.decorator_list: + if (isinstance(d, ast.Name) and d.id == "property") or ( + isinstance(d, ast.Attribute) and d.attr == "property" + ): + is_relevant = True + break + if is_relevant: + relevant_nodes.append(item) + + if not relevant_nodes: + return None + + snippets: list[str] = [] + for fn_node in relevant_nodes: + start = fn_node.lineno + if fn_node.decorator_list: + # Compute minimum decorator lineno with an explicit loop (avoids generator/min overhead) + m = start + for d in fn_node.decorator_list: + m = min(m, d.lineno) + start = m + snippets.append("\n".join(lines[start - 1 : fn_node.end_lineno])) + + return f"class {class_name}:\n" + "\n".join(snippets) + + +def extract_parameter_type_constructors( + function_to_optimize: FunctionToOptimize, project_root_path: Path, existing_class_names: set[str] +) -> CodeStringsMarkdown: + import jedi + + try: + source = function_to_optimize.file_path.read_text(encoding="utf-8") + tree = ast.parse(source) + except Exception: + return CodeStringsMarkdown(code_strings=[]) + + func_node = None + for node in ast.walk(tree): + if ( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function_to_optimize.function_name + ): + if function_to_optimize.starting_line is not None and node.lineno != function_to_optimize.starting_line: + continue + func_node = node + break + if func_node is None: + return CodeStringsMarkdown(code_strings=[]) + + type_names: set[str] = set() + for arg in func_node.args.args + func_node.args.posonlyargs + func_node.args.kwonlyargs: + type_names |= collect_type_names_from_annotation(arg.annotation) + if func_node.args.vararg: + type_names |= collect_type_names_from_annotation(func_node.args.vararg.annotation) + if func_node.args.kwarg: + type_names |= collect_type_names_from_annotation(func_node.args.kwarg.annotation) + + type_names -= BUILTIN_AND_TYPING_NAMES + type_names -= existing_class_names + if not type_names: + return CodeStringsMarkdown(code_strings=[]) + + import_map: dict[str, str] = {} + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + for alias in node.names: + name = alias.asname if alias.asname else alias.name + import_map[name] = node.module + + code_strings: list[CodeString] = [] + module_cache: dict[Path, tuple[str, ast.Module]] = {} + + for type_name in sorted(type_names): + module_name = import_map.get(type_name) + if not module_name: + continue + try: + script_code = f"from {module_name} import {type_name}" + script = jedi.Script(script_code, project=jedi.Project(path=project_root_path)) + definitions = script.goto(1, len(f"from {module_name} import ") + len(type_name), follow_imports=True) + if not definitions: + continue + + module_path = definitions[0].module_path + if not module_path: + continue + + if module_path in module_cache: + mod_source, mod_tree = module_cache[module_path] + else: + mod_source = module_path.read_text(encoding="utf-8") + mod_tree = ast.parse(mod_source) + module_cache[module_path] = (mod_source, mod_tree) + + stub = extract_init_stub_from_class(type_name, mod_source, mod_tree) + if stub: + code_strings.append(CodeString(code=stub, file_path=module_path)) + except Exception: + logger.debug(f"Error extracting constructor stub for {type_name} from {module_name}") + continue + + return CodeStringsMarkdown(code_strings=code_strings) + + +def resolve_instance_class_name(name: str, module_tree: ast.Module) -> str | None: + for node in module_tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == name: + value = node.value + if isinstance(value, ast.Call): + func = value.func + if isinstance(func, ast.Name): + return func.id + if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + return func.value.id + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == name: + ann = node.annotation + if isinstance(ann, ast.Name): + return ann.id + if isinstance(ann, ast.Subscript) and isinstance(ann.value, ast.Name): + return ann.value.id + return None + + +def enrich_testgen_context(code_context: CodeStringsMarkdown, project_root_path: Path) -> CodeStringsMarkdown: + import jedi + + result = _parse_and_collect_imports(code_context) + if result is None: + return CodeStringsMarkdown(code_strings=[]) + tree, imported_names = result + + if not imported_names: + return CodeStringsMarkdown(code_strings=[]) + + existing_classes = collect_existing_class_names(tree) + + code_strings: list[CodeString] = [] + emitted_class_names: set[str] = set() + + # --- Step 1: Project class definitions (jedi resolution + recursive base extraction) --- + extracted_classes: set[tuple[Path, str]] = set() + module_cache: dict[Path, tuple[str, ast.Module]] = {} + + def get_module_source_and_tree(module_path: Path) -> tuple[str, ast.Module] | None: + if module_path in module_cache: + return module_cache[module_path] + try: + module_source = module_path.read_text(encoding="utf-8") + module_tree = ast.parse(module_source) + except Exception: + return None + else: + module_cache[module_path] = (module_source, module_tree) + return module_source, module_tree + + def extract_class_and_bases( + class_name: str, module_path: Path, module_source: str, module_tree: ast.Module + ) -> None: + if (module_path, class_name) in extracted_classes: + return + + class_node = None + for node in ast.walk(module_tree): + if isinstance(node, ast.ClassDef) and node.name == class_name: + class_node = node + break + + if class_node is None: + return + + for base in class_node.bases: + base_name = None + if isinstance(base, ast.Name): + base_name = base.id + elif isinstance(base, ast.Attribute): + continue + + if base_name and base_name not in existing_classes: + extract_class_and_bases(base_name, module_path, module_source, module_tree) + + if (module_path, class_name) in extracted_classes: + return + + lines = module_source.split("\n") + start_line = class_node.lineno + if class_node.decorator_list: + start_line = min(d.lineno for d in class_node.decorator_list) + class_source = "\n".join(lines[start_line - 1 : class_node.end_lineno]) + + full_source = class_source + + code_strings.append(CodeString(code=full_source, file_path=module_path)) + extracted_classes.add((module_path, class_name)) + emitted_class_names.add(class_name) + + for name, module_name in imported_names.items(): + if name in existing_classes or module_name == "__future__": + continue + try: + test_code = f"import {module_name}" + script = jedi.Script(test_code, project=jedi.Project(path=project_root_path)) + completions = script.goto(1, len(test_code)) + + if not completions: + continue + + module_path = completions[0].module_path + if not module_path: + continue + + resolved_module = module_path.resolve() + module_str = str(resolved_module) + is_project = module_str.startswith(str(project_root_path.resolve()) + os.sep) + is_third_party = "site-packages" in module_str + if not is_project and not is_third_party: + continue + + mod_result = get_module_source_and_tree(module_path) + if mod_result is None: + continue + module_source, module_tree = mod_result + + extract_class_and_bases(name, module_path, module_source, module_tree) + + if (module_path, name) not in extracted_classes: + resolved_class = resolve_instance_class_name(name, module_tree) + if resolved_class and resolved_class not in existing_classes: + extract_class_and_bases(resolved_class, module_path, module_source, module_tree) + + except Exception: + logger.debug(f"Error extracting class definition for {name} from {module_name}") + continue + + return CodeStringsMarkdown(code_strings=code_strings) + + +def resolve_classes_from_modules(candidates: set[tuple[str, str]]) -> list[tuple[type, str]]: + """Import modules and resolve candidate (class_name, module_name) pairs to class objects.""" + import importlib + import inspect + + resolved: list[tuple[type, str]] = [] + module_cache: dict[str, object] = {} + + for class_name, module_name in candidates: + try: + module = module_cache.get(module_name) + if module is None: + module = importlib.import_module(module_name) + module_cache[module_name] = module + + cls = getattr(module, class_name, None) + if cls is not None and inspect.isclass(cls): + resolved.append((cls, class_name)) + except (ImportError, ModuleNotFoundError, AttributeError): + logger.debug(f"Failed to import {module_name}.{class_name}") + + return resolved + + +MAX_TRANSITIVE_DEPTH = 5 + + +def extract_classes_from_type_hint(hint: object) -> list[type]: + """Recursively extract concrete class objects from a type annotation. + + Unwraps Optional, Union, List, Dict, Callable, Annotated, etc. + Filters out builtins and typing module types. + """ + import typing + + classes: list[type] = [] + origin = getattr(hint, "__origin__", None) + args = getattr(hint, "__args__", None) + + if origin is not None and args: + for arg in args: + classes.extend(extract_classes_from_type_hint(arg)) + elif isinstance(hint, type): + module = getattr(hint, "__module__", "") + if module not in ("builtins", "typing", "typing_extensions", "types"): + classes.append(hint) + # Handle typing.Annotated on older Pythons where __origin__ may not be set + if hasattr(typing, "get_args") and origin is None and args is None: + try: + inner_args = typing.get_args(hint) + if inner_args: + for arg in inner_args: + classes.extend(extract_classes_from_type_hint(arg)) + except Exception: + pass + + return classes + + +def resolve_transitive_type_deps(cls: type) -> list[type]: + """Find external classes referenced in cls.__init__ type annotations. + + Returns classes from site-packages that have a custom __init__. + """ + import inspect + import typing + + try: + init_method = getattr(cls, "__init__") + hints = typing.get_type_hints(init_method) + except Exception: + return [] + + deps: list[type] = [] + for param_name, hint in hints.items(): + if param_name == "return": + continue + for dep_cls in extract_classes_from_type_hint(hint): + if dep_cls is cls: + continue + init_method = getattr(dep_cls, "__init__", None) + if init_method is None or init_method is object.__init__: + continue + try: + class_file = Path(inspect.getfile(dep_cls)) + except (OSError, TypeError): + continue + if not path_belongs_to_site_packages(class_file): + continue + deps.append(dep_cls) + + return deps + + +def extract_init_stub(cls: type, class_name: str, require_site_packages: bool = True) -> CodeString | None: + """Extract a stub containing the class definition with only its __init__ method. + + Args: + cls: The class object to extract __init__ from + class_name: Name to use for the class in the stub + require_site_packages: If True, only extract from site-packages. If False, include stdlib too. + + """ + import inspect + import textwrap + + init_method = getattr(cls, "__init__", None) + if init_method is None or init_method is object.__init__: + return None + + try: + class_file = Path(inspect.getfile(cls)) + except (OSError, TypeError): + return None + + if require_site_packages and not path_belongs_to_site_packages(class_file): + return None + + try: + init_source = inspect.getsource(init_method) + init_source = textwrap.dedent(init_source) + except (OSError, TypeError): + return None + + parts = class_file.parts + if "site-packages" in parts: + idx = parts.index("site-packages") + class_file = Path(*parts[idx + 1 :]) + + class_source = f"class {class_name}:\n" + textwrap.indent(init_source, " ") + return CodeString(code=class_source, file_path=class_file) + + +def _is_project_module_cached(module_name: str, project_root_path: Path, cache: dict[str, bool]) -> bool: + cached = cache.get(module_name) + if cached is not None: + return cached + is_project = _is_project_module(module_name, project_root_path) + cache[module_name] = is_project + return is_project + + +def is_project_path(module_path: Path | None, project_root_path: Path) -> bool: + if module_path is None: + return False + # site-packages must be checked first because .venv/site-packages is under project root + if path_belongs_to_site_packages(module_path): + return False + try: + module_path.resolve().relative_to(project_root_path.resolve()) + return True + except ValueError: + return False + + +def _is_project_module(module_name: str, project_root_path: Path) -> bool: + """Check if a module is part of the project (not external/stdlib).""" + import importlib.util + + try: + spec = importlib.util.find_spec(module_name) + except (ImportError, ModuleNotFoundError, ValueError): + return False + else: + if spec is None or spec.origin is None: + return False + return is_project_path(Path(spec.origin), project_root_path) + + +def extract_imports_for_class(module_tree: ast.Module, class_node: ast.ClassDef, module_source: str) -> str: + """Extract import statements needed for a class definition. + + This extracts imports for base classes, decorators, and type annotations. + """ + needed_names: set[str] = set() + + # Get base class names + for base in class_node.bases: + if isinstance(base, ast.Name): + needed_names.add(base.id) + elif isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name): + # For things like abc.ABC, we need the module name + needed_names.add(base.value.id) + + # Get decorator names (e.g., dataclass, field) + for decorator in class_node.decorator_list: + if isinstance(decorator, ast.Name): + needed_names.add(decorator.id) + elif isinstance(decorator, ast.Call): + if isinstance(decorator.func, ast.Name): + needed_names.add(decorator.func.id) + elif isinstance(decorator.func, ast.Attribute) and isinstance(decorator.func.value, ast.Name): + needed_names.add(decorator.func.value.id) + + # Get type annotation names from class body (for dataclass fields) + for item in class_node.body: + if isinstance(item, ast.AnnAssign) and item.annotation: + collect_names_from_annotation(item.annotation, needed_names) + # Also check for field() calls which are common in dataclasses + elif isinstance(item, ast.Assign) and isinstance(item.value, ast.Call): + if isinstance(item.value.func, ast.Name): + needed_names.add(item.value.func.id) + + # Find imports that provide these names + import_lines: list[str] = [] + source_lines = module_source.split("\n") + added_imports: set[int] = set() # Track line numbers to avoid duplicates + + for node in module_tree.body: + if isinstance(node, ast.Import): + for alias in node.names: + name = alias.asname if alias.asname else alias.name.split(".")[0] + if name in needed_names and node.lineno not in added_imports: + import_lines.append(source_lines[node.lineno - 1]) + added_imports.add(node.lineno) + break + elif isinstance(node, ast.ImportFrom): + for alias in node.names: + name = alias.asname if alias.asname else alias.name + if name in needed_names and node.lineno not in added_imports: + import_lines.append(source_lines[node.lineno - 1]) + added_imports.add(node.lineno) + break + + return "\n".join(import_lines) + + +def collect_names_from_annotation(node: ast.expr, names: set[str]) -> None: + """Recursively collect type annotation names from an AST node.""" + if isinstance(node, ast.Name): + names.add(node.id) + elif isinstance(node, ast.Subscript): + collect_names_from_annotation(node.value, names) + collect_names_from_annotation(node.slice, names) + elif isinstance(node, ast.Tuple): + for elt in node.elts: + collect_names_from_annotation(elt, names) + elif isinstance(node, ast.BinOp): # For Union types with | syntax + collect_names_from_annotation(node.left, names) + collect_names_from_annotation(node.right, names) + elif isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name): + names.add(node.value.id) + + +def is_dunder_method(name: str) -> bool: + return len(name) > 4 and name.isascii() and name.startswith("__") and name.endswith("__") + + +def remove_docstring_from_body(indented_block: cst.IndentedBlock) -> cst.CSTNode: + """Removes the docstring from an indented block if it exists.""" + if not isinstance(indented_block.body[0], cst.SimpleStatementLine): + return indented_block + first_stmt = indented_block.body[0].body[0] + if isinstance(first_stmt, cst.Expr) and isinstance(first_stmt.value, cst.SimpleString): + return indented_block.with_changes(body=indented_block.body[1:]) + return indented_block + + +def parse_code_and_prune_cst( + code: str, + code_context_type: CodeContextType, + target_functions: set[str], + helpers_of_helper_functions: set[str] = set(), # noqa: B006 + remove_docstrings: bool = False, +) -> cst.Module: + """Parse and filter the code CST, returning the pruned Module.""" + module = cst.parse_module(code) + defs_with_usages = collect_top_level_defs_with_usages(module, target_functions | helpers_of_helper_functions) + + if code_context_type == CodeContextType.READ_WRITABLE: + filtered_node, found_target = prune_cst( + module, target_functions, defs_with_usages=defs_with_usages, keep_class_init=True + ) + elif code_context_type == CodeContextType.READ_ONLY: + filtered_node, found_target = prune_cst( + module, + target_functions, + helpers=helpers_of_helper_functions, + remove_docstrings=remove_docstrings, + include_target_in_output=False, + include_dunder_methods=True, + ) + elif code_context_type == CodeContextType.TESTGEN: + filtered_node, found_target = prune_cst( + module, + target_functions, + helpers=helpers_of_helper_functions, + remove_docstrings=remove_docstrings, + include_dunder_methods=True, + include_init_dunder=True, + ) + elif code_context_type == CodeContextType.HASHING: + filtered_node, found_target = prune_cst( + module, target_functions, remove_docstrings=True, exclude_init_from_targets=True + ) + else: + raise ValueError(f"Unknown code_context_type: {code_context_type}") # noqa: EM102 + + if not found_target: + raise ValueError("No target functions found in the provided code") + if filtered_node and isinstance(filtered_node, cst.Module): + return filtered_node + raise ValueError("Pruning produced no module") + + +def prune_cst( + node: cst.CSTNode, + target_functions: set[str], + prefix: str = "", + *, + defs_with_usages: dict[str, UsageInfo] | None = None, + helpers: set[str] | None = None, + remove_docstrings: bool = False, + include_target_in_output: bool = True, + exclude_init_from_targets: bool = False, + keep_class_init: bool = False, + include_dunder_methods: bool = False, + include_init_dunder: bool = False, +) -> tuple[cst.CSTNode | None, bool]: + """Unified function to prune CST nodes based on various filtering criteria. + + Args: + node: The CST node to filter + target_functions: Set of qualified function names that are targets + prefix: Current qualified name prefix (for class methods) + defs_with_usages: Dict of definitions with usage info (for READ_WRITABLE mode) + helpers: Set of helper function qualified names (for READ_ONLY/TESTGEN modes) + remove_docstrings: Whether to remove docstrings from output + include_target_in_output: Whether to include target functions in output + exclude_init_from_targets: Whether to exclude __init__ from targets (HASHING mode) + keep_class_init: Whether to keep __init__ methods in classes (READ_WRITABLE mode) + include_dunder_methods: Whether to include dunder methods (READ_ONLY/TESTGEN modes) + include_init_dunder: Whether to include __init__ in dunder methods + + Returns: + (filtered_node, found_target): + filtered_node: The modified CST node or None if it should be removed. + found_target: True if a target function was found in this node's subtree. + + """ + if isinstance(node, (cst.Import, cst.ImportFrom)): + return None, False + + if isinstance(node, cst.FunctionDef): + qualified_name = f"{prefix}.{node.name.value}" if prefix else node.name.value + + # Check if it's a helper function (higher priority than target) + if helpers and qualified_name in helpers: + if remove_docstrings and isinstance(node.body, cst.IndentedBlock): + return node.with_changes(body=remove_docstring_from_body(node.body)), True + return node, True + + # Check if it's a target function + if qualified_name in target_functions: + # Handle exclude_init_from_targets for HASHING mode + if exclude_init_from_targets and node.name.value == "__init__": + return None, False + + if include_target_in_output: + if remove_docstrings and isinstance(node.body, cst.IndentedBlock): + return node.with_changes(body=remove_docstring_from_body(node.body)), True + return node, True + return None, True + + # Handle class __init__ for READ_WRITABLE mode + if keep_class_init and node.name.value == "__init__": + return node, False + + # Handle dunder methods for READ_ONLY/TESTGEN modes + if ( + include_dunder_methods + and len(node.name.value) > 4 + and node.name.value.startswith("__") + and node.name.value.endswith("__") + ): + if not include_init_dunder and node.name.value == "__init__": + return None, False + if remove_docstrings and isinstance(node.body, cst.IndentedBlock): + return node.with_changes(body=remove_docstring_from_body(node.body)), False + return node, False + + return None, False + + if isinstance(node, cst.ClassDef): + if prefix: + return None, False + if not isinstance(node.body, cst.IndentedBlock): + raise ValueError("ClassDef body is not an IndentedBlock") # noqa: TRY004 + class_prefix = node.name.value + class_name = node.name.value + + # Handle dependency classes for READ_WRITABLE mode + if defs_with_usages: + # Check if this class contains any target functions + has_target_functions = any( + isinstance(stmt, cst.FunctionDef) and f"{class_prefix}.{stmt.name.value}" in target_functions + for stmt in node.body.body + ) + + # If the class is used as a dependency (not containing target functions), keep it entirely + if ( + not has_target_functions + and class_name in defs_with_usages + and defs_with_usages[class_name].used_by_qualified_function + ): + return node, True + + # Recursively filter each statement in the class body + new_class_body: list[cst.CSTNode] = [] + found_in_class = False + + for stmt in node.body.body: + filtered, found_target = prune_cst( + stmt, + target_functions, + class_prefix, + defs_with_usages=defs_with_usages, + helpers=helpers, + remove_docstrings=remove_docstrings, + include_target_in_output=include_target_in_output, + exclude_init_from_targets=exclude_init_from_targets, + keep_class_init=keep_class_init, + include_dunder_methods=include_dunder_methods, + include_init_dunder=include_init_dunder, + ) + found_in_class |= found_target + if filtered: + new_class_body.append(filtered) + + if not found_in_class: + return None, False + + # Apply docstring removal to class if needed + if remove_docstrings and new_class_body: + updated_body = node.body.with_changes(body=new_class_body) + assert isinstance(updated_body, cst.IndentedBlock) + return node.with_changes(body=remove_docstring_from_body(updated_body)), True + + return node.with_changes(body=node.body.with_changes(body=new_class_body)) if new_class_body else None, True + + # Handle assignments for READ_WRITABLE mode + if defs_with_usages is not None: + 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 + section_names = get_section_names(node) + if not section_names: + return node, False + + if helpers is not None: + return recurse_sections( + node, + section_names, + lambda child: prune_cst( + child, + target_functions, + prefix, + defs_with_usages=defs_with_usages, + helpers=helpers, + remove_docstrings=remove_docstrings, + include_target_in_output=include_target_in_output, + exclude_init_from_targets=exclude_init_from_targets, + keep_class_init=keep_class_init, + include_dunder_methods=include_dunder_methods, + include_init_dunder=include_init_dunder, + ), + keep_non_target_children=True, + ) + return recurse_sections( + node, + section_names, + lambda child: prune_cst( + child, + target_functions, + prefix, + defs_with_usages=defs_with_usages, + helpers=helpers, + remove_docstrings=remove_docstrings, + include_target_in_output=include_target_in_output, + exclude_init_from_targets=exclude_init_from_targets, + keep_class_init=keep_class_init, + include_dunder_methods=include_dunder_methods, + include_init_dunder=include_init_dunder, + ), + ) diff --git a/codeflash/languages/python/context/unused_definition_remover.py b/codeflash/languages/python/context/unused_definition_remover.py new file mode 100644 index 000000000..e70dcad29 --- /dev/null +++ b/codeflash/languages/python/context/unused_definition_remover.py @@ -0,0 +1,854 @@ +from __future__ import annotations + +import ast +from collections import defaultdict +from dataclasses import dataclass, field +from itertools import chain +from pathlib import Path +from typing import TYPE_CHECKING, Optional, Union + +import libcst as cst + +from codeflash.cli_cmds.console import logger +from codeflash.languages import is_python +from codeflash.languages.python.static_analysis.code_replacer import replace_function_definitions_in_module +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 + + +@dataclass +class UsageInfo: + """Information about a name and its usage.""" + + name: str + used_by_qualified_function: bool = False + dependencies: set[str] = field(default_factory=set) + + +def extract_names_from_targets(target: cst.CSTNode) -> list[str]: + """Extract all variable names from a target node, including from tuple unpacking.""" + names = [] + + # Handle a simple name + if isinstance(target, cst.Name): + names.append(target.value) + + # Handle any node with a value attribute (StarredElement, etc.) + elif hasattr(target, "value"): + names.extend(extract_names_from_targets(target.value)) + + # Handle any node with elements attribute (tuples, lists, etc.) + elif hasattr(target, "elements"): + for element in target.elements: + # Recursive call for each element + names.extend(extract_names_from_targets(element)) + + 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]: + """Recursively collect all top-level variable, function, and class definitions.""" + # Locally bind types and helpers for faster lookup + FunctionDef = cst.FunctionDef # noqa: N806 + ClassDef = cst.ClassDef # noqa: N806 + Assign = cst.Assign # noqa: N806 + AnnAssign = cst.AnnAssign # noqa: N806 + AugAssign = cst.AugAssign # noqa: N806 + IndentedBlock = cst.IndentedBlock # noqa: N806 + + if definitions is None: + definitions = {} + + # Speed: Single isinstance+local var instead of several type calls + node_type = type(node) + # Fast path: function def + if node_type is FunctionDef: + name = node.name.value + definitions[name] = UsageInfo( + name=name, + used_by_qualified_function=False, # Will be marked later if in qualified functions + ) + return definitions + + # Fast path: class def + if node_type is ClassDef: + name = node.name.value + definitions[name] = UsageInfo(name=name) + + # Collect class methods + body = getattr(node, "body", None) + if body is not None and type(body) is IndentedBlock: + statements = body.body + # Precompute f-string template for efficiency + prefix = name + "." + for statement in statements: + if type(statement) is FunctionDef: + method_name = prefix + statement.name.value + definitions[method_name] = UsageInfo(name=method_name) + + return definitions + + # Fast path: assignment + if node_type is Assign: + # Inline extract_names_from_targets for single-target speed + targets = node.targets + append_def = definitions.__setitem__ + for target in targets: + names = extract_names_from_targets(target.target) + for name in names: + append_def(name, UsageInfo(name=name)) + return definitions + + if node_type is AnnAssign or node_type is AugAssign: + tgt = node.target + if type(tgt) is cst.Name: + name = tgt.value + definitions[name] = UsageInfo(name=name) + else: + names = extract_names_from_targets(tgt) + for name in names: + definitions[name] = UsageInfo(name=name) + return definitions + + # Recursively process children. Takes care of top level assignments in if/else/while/for blocks + section_names = get_section_names(node) + + if section_names: + getattr_ = getattr + for section in section_names: + original_content = getattr_(node, section, None) + # Instead of isinstance check for list/tuple, rely on duck-type via iter + # If section contains a list of nodes + if isinstance(original_content, (list, tuple)): + defs = definitions # Move out for minor speed + for child in original_content: + collect_top_level_definitions(child, defs) + # If section contains a single node + elif original_content is not None: + collect_top_level_definitions(original_content, definitions) + + return definitions + + +def get_section_names(node: cst.CSTNode) -> list[str]: + """Return the section attribute names (e.g., body, orelse) for a given node if they exist.""" + possible_sections = ["body", "orelse", "finalbody", "handlers"] + return [sec for sec in possible_sections if hasattr(node, sec)] + + +class DependencyCollector(cst.CSTVisitor): + """Collects dependencies between definitions using the visitor pattern with depth tracking.""" + + METADATA_DEPENDENCIES = (cst.metadata.ParentNodeProvider,) + + def __init__(self, definitions: dict[str, UsageInfo]) -> None: + super().__init__() + self.definitions = definitions + # Track function and class depths + self.function_depth = 0 + self.class_depth = 0 + # Track top-level qualified names + self.current_top_level_name = "" + self.current_class = "" + # Track if we're processing a top-level variable + self.processing_variable = False + self.current_variable_names = set() + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + function_name = node.name.value + + if self.function_depth == 0: + # This is a top-level function + if self.class_depth > 0: + # If inside a class, we're now tracking dependencies at the class level + self.current_top_level_name = f"{self.current_class}.{function_name}" + else: + # Regular top-level function + self.current_top_level_name = function_name + + # Check parameter type annotations for dependencies + if hasattr(node, "params") and node.params: + for param in node.params.params: + if param.annotation: + # Visit the annotation to extract dependencies + self._collect_annotation_dependencies(param.annotation) + + self.function_depth += 1 + + def _collect_annotation_dependencies(self, annotation: cst.Annotation) -> None: + """Extract dependencies from type annotations.""" + if hasattr(annotation, "annotation"): + # Extract names from annotation (could be Name, Attribute, Subscript, etc.) + self._extract_names_from_annotation(annotation.annotation) + + def _extract_names_from_annotation(self, node: cst.CSTNode) -> None: + """Extract names from a type annotation node.""" + # Simple name reference like 'int', 'str', or custom type + if isinstance(node, cst.Name): + name = node.value + if name in self.definitions and name != self.current_top_level_name and self.current_top_level_name: + self.definitions[self.current_top_level_name].dependencies.add(name) + + # Handle compound annotations like List[int], Dict[str, CustomType], etc. + elif isinstance(node, cst.Subscript): + if hasattr(node, "value"): + self._extract_names_from_annotation(node.value) + if hasattr(node, "slice"): + for slice_item in node.slice: + if hasattr(slice_item, "slice"): + self._extract_names_from_annotation(slice_item.slice) + + # Handle attribute access like module.Type + elif isinstance(node, cst.Attribute): + if hasattr(node, "value"): + self._extract_names_from_annotation(node.value) + # No need to check the attribute name itself as it's likely not a top-level definition + + def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: + self.function_depth -= 1 + + if self.function_depth == 0 and self.class_depth == 0: + # Exiting top-level function that's not in a class + self.current_top_level_name = "" + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + class_name = node.name.value + + if self.class_depth == 0: + # This is a top-level class + self.current_class = class_name + self.current_top_level_name = class_name + + # Track base classes as dependencies + for base in node.bases: + if isinstance(base.value, cst.Name): + base_name = base.value.value + if base_name in self.definitions and class_name in self.definitions: + self.definitions[class_name].dependencies.add(base_name) + elif isinstance(base.value, cst.Attribute): + # Handle cases like module.ClassName + attr_name = base.value.attr.value + if attr_name in self.definitions and class_name in self.definitions: + self.definitions[class_name].dependencies.add(attr_name) + + self.class_depth += 1 + + def leave_ClassDef(self, original_node: cst.ClassDef) -> None: + self.class_depth -= 1 + + if self.class_depth == 0: + # Exiting top-level class + self.current_class = "" + self.current_top_level_name = "" + + def visit_Assign(self, node: cst.Assign) -> None: + # Only handle top-level assignments + if self.function_depth == 0 and self.class_depth == 0: + for target in node.targets: + # Extract all variable names from the target + names = extract_names_from_targets(target.target) + + # Check if any of these names are top-level definitions we're tracking + tracked_names = [name for name in names if name in self.definitions] + if tracked_names: + self.processing_variable = True + self.current_variable_names.update(tracked_names) + # Use the first tracked name as the current top-level name (for dependency tracking) + self.current_top_level_name = tracked_names[0] + + def leave_Assign(self, original_node: cst.Assign) -> None: + if self.processing_variable: + self.processing_variable = False + self.current_variable_names.clear() + self.current_top_level_name = "" + + def visit_AnnAssign(self, node: cst.AnnAssign) -> None: + # Extract names from the variable annotations + if hasattr(node, "annotation") and node.annotation: + # First mark we're processing a variable to avoid recording it as a dependency of itself + self.processing_variable = True + if isinstance(node.target, cst.Name): + self.current_variable_names.add(node.target.value) + else: + self.current_variable_names.update(extract_names_from_targets(node.target)) + + # Process the annotation + self._collect_annotation_dependencies(node.annotation) + + # Reset processing state + self.processing_variable = False + self.current_variable_names.clear() + + def visit_Name(self, node: cst.Name) -> None: + name = node.value + + # Skip if we're not inside a tracked definition + if not self.current_top_level_name or self.current_top_level_name not in self.definitions: + return + + # Skip if we're looking at the variable name itself in an assignment + if self.processing_variable and name in self.current_variable_names: + return + + if name in self.definitions and name != self.current_top_level_name: + # Skip if this Name is the .attr part of an Attribute (e.g., 'x' in 'self.x') + # We only want to track the base/value of attribute access, not the attribute name itself + if self.class_depth > 0: + parent = self.get_metadata(cst.metadata.ParentNodeProvider, node) + if parent is not None and isinstance(parent, cst.Attribute): + # Check if this Name is the .attr (property name), not the .value (base) + # If it's the .attr, skip it - attribute names aren't references to definitions + if parent.attr is node: + return + # If it's the .value (base), only skip if it's self/cls + if name in ("self", "cls"): + return + self.definitions[self.current_top_level_name].dependencies.add(name) + + +class QualifiedFunctionUsageMarker: + """Marks definitions that are used by specific qualified functions.""" + + def __init__(self, definitions: dict[str, UsageInfo], qualified_function_names: set[str]) -> None: + self.definitions = definitions + self.qualified_function_names = qualified_function_names + self.expanded_qualified_functions = self._expand_qualified_functions() + + def _expand_qualified_functions(self) -> set[str]: + """Expand the qualified function names to include related methods.""" + expanded = set(self.qualified_function_names) + + # Find class methods and add their containing classes and dunder methods + for qualified_name in list(self.qualified_function_names): + if "." in qualified_name: + class_name, _method_name = qualified_name.split(".", 1) + + # Add the class itself + expanded.add(class_name) + + # Add all dunder methods of the class + for name in self.definitions: + if name.startswith(f"{class_name}.__") and name.endswith("__"): + expanded.add(name) + + return expanded + + def mark_used_definitions(self) -> None: + """Find all qualified functions and mark them and their dependencies as used.""" + # Avoid list comprehension for set intersection + expanded_names = self.expanded_qualified_functions + defs = self.definitions + # Use set intersection but only if defs.keys is a set (Python 3.12 dict_keys supports it efficiently) + fnames = ( + expanded_names & defs.keys() + if isinstance(expanded_names, set) + else [name for name in expanded_names if name in defs] + ) + + # For each specified function, mark it and all its dependencies as used + for func_name in fnames: + defs[func_name].used_by_qualified_function = True + for dep in defs[func_name].dependencies: + self.mark_as_used_recursively(dep) + + def mark_as_used_recursively(self, name: str) -> None: + """Mark a name and all its dependencies as used recursively.""" + if name not in self.definitions: + return + + if self.definitions[name].used_by_qualified_function: + return # Already marked + + self.definitions[name].used_by_qualified_function = True + + # Mark all dependencies as used + for dep in self.definitions[name].dependencies: + self.mark_as_used_recursively(dep) + + +def remove_unused_definitions_recursively( + node: cst.CSTNode, definitions: dict[str, UsageInfo] +) -> tuple[cst.CSTNode | None, bool]: + """Recursively filter the node to remove unused definitions. + + Args: + ---- + node: The CST node to process + definitions: Dictionary of definition info + + Returns: + ------- + (filtered_node, used_by_function): + filtered_node: The modified CST node or None if it should be removed + used_by_function: True if this node or any child is used by qualified functions + + """ + # Skip import statements + if isinstance(node, (cst.Import, cst.ImportFrom)): + return node, True + + # Never remove function definitions + if isinstance(node, cst.FunctionDef): + return node, True + + # Never remove class definitions + if isinstance(node, cst.ClassDef): + class_name = node.name.value + + # Check if any methods or variables in this class are used + method_or_var_used = False + class_has_dependencies = False + + # Check if class itself is marked as used + if class_name in definitions and definitions[class_name].used_by_qualified_function: + class_has_dependencies = True + + if hasattr(node, "body") and isinstance(node.body, cst.IndentedBlock): + updates = {} + new_statements = [] + + for statement in node.body.body: + # Keep all function definitions + if isinstance(statement, cst.FunctionDef): + method_name = f"{class_name}.{statement.name.value}" + if method_name in definitions and definitions[method_name].used_by_qualified_function: + method_or_var_used = True + new_statements.append(statement) + # Only process variable assignments + elif isinstance(statement, (cst.Assign, cst.AnnAssign, cst.AugAssign)): + var_used = False + + 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) + else: + # Keep all other statements in the class + new_statements.append(statement) + + # Update the class body + new_body = node.body.with_changes(body=new_statements) + updates["body"] = new_body + + return node.with_changes(**updates), True + + return node, method_or_var_used or class_has_dependencies + + # 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 + return recurse_sections( + node, section_names, lambda child: remove_unused_definitions_recursively(child, definitions) + ) + + +def collect_top_level_defs_with_usages( + code: Union[str, cst.Module], qualified_function_names: set[str] +) -> dict[str, UsageInfo]: + """Collect all top level definitions (classes, variables or functions) and their usages.""" + module = code if isinstance(code, cst.Module) else cst.parse_module(code) + # Collect all definitions (top level classes, variables or function) + definitions = collect_top_level_definitions(module) + + # Collect dependencies between definitions using the visitor pattern + wrapper = cst.MetadataWrapper(module) + dependency_collector = DependencyCollector(definitions) + wrapper.visit(dependency_collector) + + # Mark definitions used by specified functions, and their dependencies recursively + usage_marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names) + usage_marker.mark_used_definitions() + return definitions + + +def remove_unused_definitions_by_function_names(code: str, qualified_function_names: set[str]) -> str: + """Analyze a file and remove top level definitions not used by specified functions. + + Top level definitions, in this context, are only classes, variables or functions. + If a class is referenced by a qualified function, we keep the entire class. + + Args: + ---- + code: The code to process + qualified_function_names: Set of function names to keep. For methods, use format 'classname.methodname' + + """ + try: + module = cst.parse_module(code) + except Exception as e: + logger.debug(f"Failed to parse code with libcst: {type(e).__name__}: {e}") + return code + + try: + defs_with_usages = collect_top_level_defs_with_usages(module, qualified_function_names) + + # Apply the recursive removal transformation + modified_module, _ = remove_unused_definitions_recursively(module, defs_with_usages) + + return modified_module.code if modified_module else "" + except Exception as e: + # If any other error occurs during processing, return the original code + logger.debug(f"Error processing code to remove unused definitions: {type(e).__name__}: {e}") + return code + + +def revert_unused_helper_functions( + project_root: Path, unused_helpers: list[FunctionSource], original_helper_code: dict[Path, str] +) -> None: + """Revert unused helper functions back to their original definitions. + + Args: + project_root: project_root + unused_helpers: List of unused helper functions to revert + original_helper_code: Dictionary mapping file paths to their original code + + """ + if not unused_helpers: + return + + logger.debug(f"Reverting {len(unused_helpers)} unused helper function(s) to original definitions") + + # Resolve all path keys for consistent comparison (Windows 8.3 short names may differ from Jedi-resolved paths) + resolved_original_helper_code = {p.resolve(): code for p, code in original_helper_code.items()} + + # Group unused helpers by file path + unused_helpers_by_file = defaultdict(list) + for helper in unused_helpers: + unused_helpers_by_file[helper.file_path.resolve()].append(helper) + + # For each file, revert the unused helper functions to their original definitions + for file_path, helpers_in_file in unused_helpers_by_file.items(): + if file_path in resolved_original_helper_code: + try: + # Get original code for this file + original_code = resolved_original_helper_code[file_path] + + # Use the code replacer to selectively revert only the unused helper functions + helper_names = [helper.qualified_name for helper in helpers_in_file] + reverted_code = replace_function_definitions_in_module( + function_names=helper_names, + optimized_code=CodeStringsMarkdown( + code_strings=[ + CodeString(code=original_code, file_path=Path(file_path).relative_to(project_root)) + ] + ), # Use original code as the "optimized" code to revert + module_abspath=file_path, + preexisting_objects=set(), # Empty set since we're reverting + project_root_path=project_root, + should_add_global_assignments=False, # since we revert helpers functions after applying the optimization, we know that the file already has global assignments added, otherwise they would be added twice. + ) + + if reverted_code: + logger.debug(f"Reverted unused helpers in {file_path}: {', '.join(helper_names)}") + + except Exception as e: + logger.error(f"Error reverting unused helpers in {file_path}: {e}") + + +def _analyze_imports_in_optimized_code( + optimized_ast: ast.AST, code_context: CodeOptimizationContext +) -> dict[str, set[str]]: + """Analyze import statements in optimized code to map imported names to qualified helper names. + + Args: + optimized_ast: The AST of the optimized code + code_context: The code optimization context containing helper functions + + Returns: + Dictionary mapping imported names to sets of possible qualified helper names + + """ + imported_names_map = defaultdict(set) + + # Precompute a two-level dict: module_name -> func_name -> [helpers] + helpers_by_file_and_func = defaultdict(dict) + helpers_by_file = defaultdict(list) # preserved for "import module" + for helper in code_context.helper_functions: + jedi_type = helper.definition_type + if jedi_type != "class": # Include when definition_type is None (non-Python) + func_name = helper.only_function_name + module_name = helper.file_path.stem + # Cache function lookup for this (module, func) + helpers_by_file_and_func[module_name].setdefault(func_name, []).append(helper) + helpers_by_file[module_name].append(helper) + + # Collect only import nodes to avoid per-node isinstance checks across the whole AST + class _ImportCollector(ast.NodeVisitor): + def __init__(self) -> None: + self.nodes: list[ast.AST] = [] + + def visit_Import(self, node: ast.Import) -> None: + self.nodes.append(node) + # No need to recurse further for import nodes + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + self.nodes.append(node) + # No need to recurse further for import-from nodes + + collector = _ImportCollector() + collector.visit(optimized_ast) + + for node in collector.nodes: + if isinstance(node, ast.ImportFrom): + # Handle "from module import function" statements + module_name = node.module + if module_name: + file_entry = helpers_by_file_and_func.get(module_name) + if file_entry: + for alias in node.names: + imported_name = alias.asname if alias.asname else alias.name + original_name = alias.name + helpers = file_entry.get(original_name) + if helpers: + imported_set = imported_names_map[imported_name] + for helper in helpers: + imported_set.add(helper.qualified_name) + imported_set.add(helper.fully_qualified_name) + + elif isinstance(node, ast.Import): + # Handle "import module" statements + for alias in node.names: + imported_name = alias.asname if alias.asname else alias.name + module_name = alias.name + helpers = helpers_by_file.get(module_name) + if helpers: + imported_set = imported_names_map[f"{imported_name}.{{func}}"] + for helper in helpers: + # For "import module" statements, functions would be called as module.function + full_call = f"{imported_name}.{helper.only_function_name}" + full_call_set = imported_names_map[full_call] + full_call_set.add(helper.qualified_name) + full_call_set.add(helper.fully_qualified_name) + + return dict(imported_names_map) + + +def find_target_node( + root: ast.AST, function_to_optimize: FunctionToOptimize +) -> Optional[ast.FunctionDef | ast.AsyncFunctionDef]: + parents = function_to_optimize.parents + node = root + for parent in parents: + # Fast loop: directly look for the matching ClassDef in node.body + body = getattr(node, "body", None) + if not body: + return None + for child in body: + if isinstance(child, ast.ClassDef) and child.name == parent.name: + node = child + break + else: + return None + + # Now node is either the root or the target parent class; look for function + body = getattr(node, "body", None) + if not body: + return None + target_name = function_to_optimize.function_name + for child in body: + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name == target_name: + return child + return None + + +def detect_unused_helper_functions( + function_to_optimize: FunctionToOptimize, + code_context: CodeOptimizationContext, + optimized_code: str | CodeStringsMarkdown, +) -> list[FunctionSource]: + """Detect helper functions that are no longer called by the optimized entrypoint function. + + Args: + function_to_optimize: The function to optimize + code_context: The code optimization context containing helper functions + optimized_code: The optimized code to analyze + + Returns: + List of FunctionSource objects representing unused helper functions + + """ + # Skip this analysis for non-Python languages since we use Python's ast module + if not is_python(): + logger.debug("Skipping unused helper function detection for non-Python languages") + return [] + + if isinstance(optimized_code, CodeStringsMarkdown) and len(optimized_code.code_strings) > 0: + return list( + chain.from_iterable( + detect_unused_helper_functions(function_to_optimize, code_context, code.code) + for code in optimized_code.code_strings + ) + ) + + try: + # Parse the optimized code to analyze function calls and imports + optimized_ast = ast.parse(optimized_code) + + # Find the optimized entrypoint function + entrypoint_function_ast = find_target_node(optimized_ast, function_to_optimize) + + if not entrypoint_function_ast: + logger.debug(f"Could not find entrypoint function {function_to_optimize.function_name} in optimized code") + return [] + + # First, analyze imports to build a mapping of imported names to their original qualified names + imported_names_map = _analyze_imports_in_optimized_code(optimized_ast, code_context) + + # Extract all function calls in the entrypoint function + called_function_names = {function_to_optimize.function_name} + for node in ast.walk(entrypoint_function_ast): + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Name): + # Regular function call: function_name() + called_name = node.func.id + called_function_names.add(called_name) + # Also add the qualified name if this is an imported function + mapped_names = imported_names_map.get(called_name) + if mapped_names: + called_function_names.update(mapped_names) + elif isinstance(node.func, ast.Attribute): + # Method call: obj.method() or self.method() or module.function() + if isinstance(node.func.value, ast.Name): + attr_name = node.func.attr + value_id = node.func.value.id + if value_id == "self": + # self.method_name() -> add both method_name and ClassName.method_name + called_function_names.add(attr_name) + # For class methods, also add the qualified name + # For class methods, also add the qualified name + if hasattr(function_to_optimize, "parents") and function_to_optimize.parents: + class_name = function_to_optimize.parents[0].name + called_function_names.add(f"{class_name}.{attr_name}") + else: + called_function_names.add(attr_name) + full_call = f"{value_id}.{attr_name}" + called_function_names.add(full_call) + # Check if this is a module.function call that maps to a helper + mapped_names = imported_names_map.get(full_call) + if mapped_names: + called_function_names.update(mapped_names) + # Handle nested attribute access like obj.attr.method() + # Handle nested attribute access like obj.attr.method() + else: + called_function_names.add(node.func.attr) + + logger.debug(f"Functions called in optimized entrypoint: {called_function_names}") + logger.debug(f"Imported names mapping: {imported_names_map}") + + # Find helper functions that are no longer called + unused_helpers = [] + entrypoint_file_path = function_to_optimize.file_path + for helper_function in code_context.helper_functions: + jedi_type = helper_function.definition_type + if jedi_type != "class": # Include when definition_type is None (non-Python) + # Check if the helper function is called using multiple name variants + helper_qualified_name = helper_function.qualified_name + helper_simple_name = helper_function.only_function_name + helper_fully_qualified_name = helper_function.fully_qualified_name + + # Check membership efficiently - exit early on first match + if ( + helper_qualified_name in called_function_names + or helper_simple_name in called_function_names + or helper_fully_qualified_name in called_function_names + ): + is_called = True + # For cross-file helpers, also consider module-based calls + elif helper_function.file_path != entrypoint_file_path: + # Add potential module.function combinations + module_name = helper_function.file_path.stem + module_call = f"{module_name}.{helper_simple_name}" + is_called = module_call in called_function_names + else: + is_called = False + + if not is_called: + unused_helpers.append(helper_function) + logger.debug(f"Helper function {helper_qualified_name} is not called in optimized code") + else: + logger.debug(f"Helper function {helper_qualified_name} is still called in optimized code") + + except Exception as e: + logger.debug(f"Error detecting unused helper functions: {e}") + return [] + else: + return unused_helpers diff --git a/codeflash/languages/python/reference_graph.py b/codeflash/languages/python/reference_graph.py new file mode 100644 index 000000000..4f389fd66 --- /dev/null +++ b/codeflash/languages/python/reference_graph.py @@ -0,0 +1,544 @@ +from __future__ import annotations + +import hashlib +import os +import sqlite3 +from collections import defaultdict +from pathlib import Path +from typing import TYPE_CHECKING + +from codeflash.cli_cmds.console import logger +from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages +from codeflash.languages.base import IndexResult +from codeflash.models.models import FunctionSource + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + from jedi.api.classes import Name + + +# --------------------------------------------------------------------------- +# Module-level helpers (must be top-level for ProcessPoolExecutor pickling) +# --------------------------------------------------------------------------- +# TODO: create call graph. + +_PARALLEL_THRESHOLD = 8 + +# Per-worker state, initialised by _init_index_worker in child processes +_worker_jedi_project: object | None = None +_worker_project_root_str: str | None = None + + +def _init_index_worker(project_root: str) -> None: + import jedi + + global _worker_jedi_project, _worker_project_root_str + _worker_jedi_project = jedi.Project(path=project_root) + _worker_project_root_str = project_root + + +def _resolve_definitions(ref: Name) -> list[Name]: + try: + inferred = ref.infer() + valid = [d for d in inferred if d.type in ("function", "class")] + if valid: + return valid + except Exception: + pass + + try: + result: list[Name] = ref.goto(follow_imports=True, follow_builtin_imports=False) + return result + except Exception: + return [] + + +def _is_valid_definition(definition: Name, caller_qualified_name: str, project_root_str: str) -> bool: + definition_path = definition.module_path + if definition_path is None: + return False + + if not str(definition_path).startswith(project_root_str + os.sep): + return False + + if path_belongs_to_site_packages(definition_path): + return False + + if not definition.full_name or not definition.full_name.startswith(definition.module_name): + return False + + if definition.type not in ("function", "class"): + return False + + try: + def_qn = get_qualified_name(definition.module_name, definition.full_name) + if def_qn == caller_qualified_name: + return False + except ValueError: + return False + + try: + from codeflash.optimization.function_context import belongs_to_function_qualified + + if belongs_to_function_qualified(definition, caller_qualified_name): + return False + except Exception: + pass + + return True + + +def _get_enclosing_function_qn(ref: Name) -> str | None: + try: + parent = ref.parent() + if parent is None or parent.type != "function": + return None + if not parent.full_name or not parent.full_name.startswith(parent.module_name): + return None + return get_qualified_name(parent.module_name, parent.full_name) + except (ValueError, AttributeError): + return None + + +def _analyze_file(file_path: Path, jedi_project: object, project_root_str: str) -> tuple[set[tuple[str, ...]], bool]: + """Pure Jedi analysis — no DB access. Returns (edges, had_error).""" + import jedi + + resolved = str(file_path.resolve()) + + try: + script = jedi.Script(path=file_path, project=jedi_project) + refs = script.get_names(all_scopes=True, definitions=False, references=True) + except Exception: + return set(), True + + edges: set[tuple[str, ...]] = set() + + for ref in refs: + try: + caller_qn = _get_enclosing_function_qn(ref) + if caller_qn is None: + continue + + definitions = _resolve_definitions(ref) + if not definitions: + continue + + definition = definitions[0] + definition_path = definition.module_path + if definition_path is None: + continue + + if not _is_valid_definition(definition, caller_qn, project_root_str): + continue + + edge_base = (resolved, caller_qn, str(definition_path)) + + if definition.type == "function": + callee_qn = get_qualified_name(definition.module_name, definition.full_name) + if len(callee_qn.split(".")) > 2: + continue + edges.add( + ( + *edge_base, + callee_qn, + definition.full_name, + definition.name, + definition.type, + definition.get_line_code(), + ) + ) + elif definition.type == "class": + init_qn = get_qualified_name(definition.module_name, f"{definition.full_name}.__init__") + if len(init_qn.split(".")) > 2: + continue + edges.add( + ( + *edge_base, + init_qn, + f"{definition.full_name}.__init__", + "__init__", + definition.type, + definition.get_line_code(), + ) + ) + except Exception: + continue + + return edges, False + + +def _index_file_worker(args: tuple[str, str]) -> tuple[str, str, set[tuple[str, ...]], bool]: + """Worker entry point for ProcessPoolExecutor.""" + file_path_str, file_hash = args + assert _worker_project_root_str is not None + edges, had_error = _analyze_file(Path(file_path_str), _worker_jedi_project, _worker_project_root_str) + return file_path_str, file_hash, edges, had_error + + +# --------------------------------------------------------------------------- + + +class ReferenceGraph: + SCHEMA_VERSION = 2 + + def __init__(self, project_root: Path, language: str = "python", db_path: Path | None = None) -> None: + import jedi + + self.project_root = project_root.resolve() + self.project_root_str = str(self.project_root) + self.language = language + self.jedi_project = jedi.Project(path=self.project_root) + + if db_path is None: + from codeflash.code_utils.compat import codeflash_cache_db + + db_path = codeflash_cache_db + + self.conn = sqlite3.connect(str(db_path)) + self.conn.execute("PRAGMA journal_mode=WAL") + self.indexed_file_hashes: dict[str, str] = {} + self._init_schema() + + def _init_schema(self) -> None: + cur = self.conn.cursor() + cur.execute("CREATE TABLE IF NOT EXISTS cg_schema_version (version INTEGER PRIMARY KEY)") + + row = cur.execute("SELECT version FROM cg_schema_version LIMIT 1").fetchone() + if row is None: + cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) + elif row[0] != self.SCHEMA_VERSION: + for table in [ + "cg_call_edges", + "cg_indexed_files", + "cg_languages", + "cg_projects", + "cg_project_meta", + "indexed_files", + "call_edges", + ]: + cur.execute(f"DROP TABLE IF EXISTS {table}") + cur.execute("DELETE FROM cg_schema_version") + cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) + + cur.execute( + """ + CREATE TABLE IF NOT EXISTS indexed_files ( + project_root TEXT NOT NULL, + language TEXT NOT NULL, + file_path TEXT NOT NULL, + file_hash TEXT NOT NULL, + PRIMARY KEY (project_root, language, file_path) + ) + """ + ) + cur.execute( + """ + CREATE TABLE IF NOT EXISTS call_edges ( + project_root TEXT NOT NULL, + language TEXT NOT NULL, + caller_file TEXT NOT NULL, + caller_qualified_name TEXT NOT NULL, + callee_file TEXT NOT NULL, + callee_qualified_name TEXT NOT NULL, + callee_fully_qualified_name TEXT NOT NULL, + callee_only_function_name TEXT NOT NULL, + callee_definition_type TEXT NOT NULL, + callee_source_line TEXT NOT NULL, + PRIMARY KEY (project_root, language, caller_file, caller_qualified_name, + callee_file, callee_qualified_name) + ) + """ + ) + cur.execute( + """ + CREATE INDEX IF NOT EXISTS idx_call_edges_caller + ON call_edges (project_root, language, caller_file, caller_qualified_name) + """ + ) + self.conn.commit() + + def get_callees( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: + file_path_to_function_source: dict[Path, set[FunctionSource]] = defaultdict(set) + function_source_list: list[FunctionSource] = [] + + all_caller_keys: list[tuple[str, str]] = [] + for file_path, qualified_names in file_path_to_qualified_names.items(): + resolved = str(file_path.resolve()) + self.ensure_file_indexed(file_path, resolved) + all_caller_keys.extend((resolved, qn) for qn in qualified_names) + + if not all_caller_keys: + return file_path_to_function_source, function_source_list + + cur = self.conn.cursor() + cur.execute("CREATE TEMP TABLE IF NOT EXISTS _caller_keys (caller_file TEXT, caller_qualified_name TEXT)") + cur.execute("DELETE FROM _caller_keys") + cur.executemany("INSERT INTO _caller_keys VALUES (?, ?)", all_caller_keys) + + rows = cur.execute( + """ + SELECT ce.callee_file, ce.callee_qualified_name, ce.callee_fully_qualified_name, + ce.callee_only_function_name, ce.callee_definition_type, ce.callee_source_line + FROM call_edges ce + INNER JOIN _caller_keys ck + ON ce.caller_file = ck.caller_file AND ce.caller_qualified_name = ck.caller_qualified_name + WHERE ce.project_root = ? AND ce.language = ? + """, + (self.project_root_str, self.language), + ).fetchall() + + for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: + callee_path = Path(callee_file) + fs = FunctionSource( + file_path=callee_path, + qualified_name=callee_qn, + fully_qualified_name=callee_fqn, + only_function_name=callee_name, + source_code=callee_src, + definition_type=callee_type, + ) + file_path_to_function_source[callee_path].add(fs) + function_source_list.append(fs) + + return file_path_to_function_source, function_source_list + + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[tuple[Path, str], int]: + all_caller_keys: list[tuple[Path, str, str]] = [] + for file_path, qualified_names in file_path_to_qualified_names.items(): + resolved = str(file_path.resolve()) + self.ensure_file_indexed(file_path, resolved) + all_caller_keys.extend((file_path, resolved, qn) for qn in qualified_names) + + if not all_caller_keys: + return {} + + cur = self.conn.cursor() + cur.execute("CREATE TEMP TABLE IF NOT EXISTS _count_keys (caller_file TEXT, caller_qualified_name TEXT)") + cur.execute("DELETE FROM _count_keys") + cur.executemany( + "INSERT INTO _count_keys VALUES (?, ?)", [(resolved, qn) for _, resolved, qn in all_caller_keys] + ) + + rows = cur.execute( + """ + SELECT ck.caller_file, ck.caller_qualified_name, COUNT(ce.rowid) + FROM _count_keys ck + LEFT JOIN call_edges ce + ON ce.caller_file = ck.caller_file AND ce.caller_qualified_name = ck.caller_qualified_name + AND ce.project_root = ? AND ce.language = ? + GROUP BY ck.caller_file, ck.caller_qualified_name + """, + (self.project_root_str, self.language), + ).fetchall() + + resolved_to_path: dict[str, Path] = {resolved: fp for fp, resolved, _ in all_caller_keys} + counts: dict[tuple[Path, str], int] = {} + for caller_file, caller_qn, cnt in rows: + counts[(resolved_to_path[caller_file], caller_qn)] = cnt + + return counts + + def ensure_file_indexed(self, file_path: Path, resolved: str | None = None) -> IndexResult: + if resolved is None: + resolved = str(file_path.resolve()) + + # Always read and hash the file before checking the cache so we detect on-disk changes + try: + content = file_path.read_text(encoding="utf-8") + except Exception: + return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True) + + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + + if self._is_file_cached(resolved, file_hash): + return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) + + return self.index_file(file_path, file_hash, resolved) + + def index_file(self, file_path: Path, file_hash: str, resolved: str | None = None) -> IndexResult: + if resolved is None: + resolved = str(file_path.resolve()) + edges, had_error = _analyze_file(file_path, self.jedi_project, self.project_root_str) + if had_error: + logger.debug(f"ReferenceGraph: failed to parse {file_path}") + return self._persist_edges(file_path, resolved, file_hash, edges, had_error) + + def _persist_edges( + self, file_path: Path, resolved: str, file_hash: str, edges: set[tuple[str, ...]], had_error: bool + ) -> IndexResult: + cur = self.conn.cursor() + scope = (self.project_root_str, self.language) + + # Clear existing data for this file + cur.execute( + "DELETE FROM call_edges WHERE project_root = ? AND language = ? AND caller_file = ?", (*scope, resolved) + ) + cur.execute( + "DELETE FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", (*scope, resolved) + ) + + # Insert new edges if parsing succeeded + if not had_error and edges: + cur.executemany( + """ + INSERT OR REPLACE INTO call_edges + (project_root, language, caller_file, caller_qualified_name, + callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + [(*scope, *edge) for edge in edges], + ) + + # Record that this file has been indexed + cur.execute( + "INSERT OR REPLACE INTO indexed_files (project_root, language, file_path, file_hash) VALUES (?, ?, ?, ?)", + (*scope, resolved, file_hash), + ) + + self.conn.commit() + self.indexed_file_hashes[resolved] = file_hash + + # Build summary for return value + edges_summary = tuple( + (caller_qn, callee_name, caller_file != callee_file) + for (caller_file, caller_qn, callee_file, _, _, callee_name, _, _) in edges + ) + cross_file_count = sum(is_cross_file for _, _, is_cross_file in edges_summary) + + return IndexResult( + file_path=file_path, + cached=False, + num_edges=len(edges), + edges=edges_summary, + cross_file_edges=cross_file_count, + error=had_error, + ) + + def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: + """Pre-index a batch of files, using multiprocessing for large uncached batches.""" + to_index: list[tuple[Path, str, str]] = [] + + for file_path in file_paths: + resolved = str(file_path.resolve()) + + try: + content = file_path.read_text(encoding="utf-8") + except Exception: + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True + ), + ) + continue + + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + + # Check if already cached (in-memory or DB) + if self._is_file_cached(resolved, file_hash): + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False + ), + ) + continue + + to_index.append((file_path, resolved, file_hash)) + + if not to_index: + return + + # Index uncached files + if len(to_index) >= _PARALLEL_THRESHOLD: + self._build_index_parallel(to_index, on_progress) + else: + for file_path, resolved, file_hash in to_index: + result = self.index_file(file_path, file_hash, resolved) + self._report_progress(on_progress, result) + + def _is_file_cached(self, resolved: str, file_hash: str) -> bool: + """Check if file is cached in memory or DB.""" + if self.indexed_file_hashes.get(resolved) == file_hash: + return True + + row = self.conn.execute( + "SELECT file_hash FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", + (self.project_root_str, self.language, resolved), + ).fetchone() + + if row and row[0] == file_hash: + self.indexed_file_hashes[resolved] = file_hash + return True + + return False + + def _report_progress(self, on_progress: Callable[[IndexResult], None] | None, result: IndexResult) -> None: + """Report progress if callback provided.""" + if on_progress is not None: + on_progress(result) + + def _build_index_parallel( + self, to_index: list[tuple[Path, str, str]], on_progress: Callable[[IndexResult], None] | None + ) -> None: + from concurrent.futures import ProcessPoolExecutor, as_completed + + max_workers = min(os.cpu_count() or 1, len(to_index), 8) + path_info: dict[str, tuple[Path, str]] = {resolved: (fp, fh) for fp, resolved, fh in to_index} + worker_args = [(resolved, fh) for _fp, resolved, fh in to_index] + + logger.debug(f"ReferenceGraph: indexing {len(to_index)} files across {max_workers} workers") + + try: + with ProcessPoolExecutor( + max_workers=max_workers, initializer=_init_index_worker, initargs=(self.project_root_str,) + ) as executor: + futures = {executor.submit(_index_file_worker, args): args[0] for args in worker_args} + + for future in as_completed(futures): + resolved = futures[future] + file_path, file_hash = path_info[resolved] + + try: + _, _, edges, had_error = future.result() + except Exception: + logger.debug(f"ReferenceGraph: worker failed for {file_path}") + self._persist_edges(file_path, resolved, file_hash, set(), had_error=True) + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True + ), + ) + continue + + if had_error: + logger.debug(f"ReferenceGraph: failed to parse {file_path}") + + result = self._persist_edges(file_path, resolved, file_hash, edges, had_error) + self._report_progress(on_progress, result) + + except Exception: + logger.debug("ReferenceGraph: parallel indexing failed, falling back to sequential") + self._fallback_sequential_index(to_index, on_progress) + + def _fallback_sequential_index( + self, to_index: list[tuple[Path, str, str]], on_progress: Callable[[IndexResult], None] | None + ) -> None: + """Fallback to sequential indexing when parallel processing fails.""" + for file_path, resolved, file_hash in to_index: + # Skip files already persisted before the failure + if resolved in self.indexed_file_hashes: + continue + result = self.index_file(file_path, file_hash, resolved) + self._report_progress(on_progress, result) + + def close(self) -> None: + self.conn.close() diff --git a/codeflash/languages/python/static_analysis/__init__.py b/codeflash/languages/python/static_analysis/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/languages/python/static_analysis/code_extractor.py b/codeflash/languages/python/static_analysis/code_extractor.py new file mode 100644 index 000000000..49b9e2b1b --- /dev/null +++ b/codeflash/languages/python/static_analysis/code_extractor.py @@ -0,0 +1,1737 @@ +from __future__ import annotations + +import ast +import time +from dataclasses import dataclass +from importlib.util import find_spec +from itertools import chain +from pathlib import Path +from typing import TYPE_CHECKING, Optional, Union + +import jedi +import libcst as cst +from libcst.codemod import CodemodContext +from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor +from libcst.helpers import calculate_module_and_package + +from codeflash.cli_cmds.console import logger +from codeflash.code_utils.config_consts import MAX_CONTEXT_LEN_REVIEW +from codeflash.languages.base import Language +from codeflash.models.models import CodePosition, FunctionParent + +if TYPE_CHECKING: + from libcst.helpers import ModuleNameAndPackage + + from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.models.models import FunctionSource + + +class GlobalFunctionCollector(cst.CSTVisitor): + """Collects all module-level function definitions (not inside classes or other functions).""" + + def __init__(self) -> None: + super().__init__() + self.functions: dict[str, cst.FunctionDef] = {} + self.function_order: list[str] = [] + self.scope_depth = 0 + + def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]: + if self.scope_depth == 0: + # Module-level function + name = node.name.value + self.functions[name] = node + if name not in self.function_order: + self.function_order.append(name) + self.scope_depth += 1 + return True + + def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: + self.scope_depth -= 1 + + def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]: + self.scope_depth += 1 + return True + + def leave_ClassDef(self, original_node: cst.ClassDef) -> None: + self.scope_depth -= 1 + + +class GlobalFunctionTransformer(cst.CSTTransformer): + """Transforms/adds module-level functions from the new file to the original file.""" + + def __init__(self, new_functions: dict[str, cst.FunctionDef], new_function_order: list[str]) -> None: + super().__init__() + self.new_functions = new_functions + self.new_function_order = new_function_order + self.processed_functions: set[str] = set() + self.scope_depth = 0 + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + self.scope_depth += 1 + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + self.scope_depth -= 1 + if self.scope_depth > 0: + return updated_node + + # Check if this is a module-level function we need to replace + name = original_node.name.value + if name in self.new_functions: + self.processed_functions.add(name) + return self.new_functions[name] + return updated_node + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + self.scope_depth += 1 + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: + self.scope_depth -= 1 + return updated_node + + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + # Add any new functions that weren't in the original file + new_statements = list(updated_node.body) + + functions_to_append = [ + self.new_functions[name] + for name in self.new_function_order + if name not in self.processed_functions and name in self.new_functions + ] + + if functions_to_append: + # Find the position of the last function or class definition + insert_index = find_insertion_index_after_imports(updated_node) + for i, stmt in enumerate(new_statements): + if isinstance(stmt, (cst.FunctionDef, cst.ClassDef)): + insert_index = i + 1 + + # Add empty line before each new function + function_nodes = [] + for func in functions_to_append: + func_with_empty_line = func.with_changes(leading_lines=[cst.EmptyLine(), *func.leading_lines]) + function_nodes.append(func_with_empty_line) + + new_statements = list(chain(new_statements[:insert_index], function_nodes, new_statements[insert_index:])) + + return updated_node.with_changes(body=new_statements) + + +def collect_referenced_names(node: cst.CSTNode) -> set[str]: + """Collect all names referenced in a CST node using recursive traversal.""" + names: set[str] = set() + + def _collect(n: cst.CSTNode) -> None: + if isinstance(n, cst.Name): + names.add(n.value) + # Recursively process all children + for child in n.children: + _collect(child) + + _collect(node) + return names + + +class GlobalAssignmentCollector(cst.CSTVisitor): + """Collects all global assignment statements.""" + + def __init__(self) -> None: + super().__init__() + self.assignments: dict[str, cst.Assign | cst.AnnAssign] = {} + self.assignment_order: list[str] = [] + # Track scope depth to identify global assignments + self.scope_depth = 0 + self.if_else_depth = 0 + + def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]: + self.scope_depth += 1 + return True + + def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: + self.scope_depth -= 1 + + def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]: + self.scope_depth += 1 + return True + + def leave_ClassDef(self, original_node: cst.ClassDef) -> None: + self.scope_depth -= 1 + + def visit_If(self, node: cst.If) -> Optional[bool]: + self.if_else_depth += 1 + return True + + def leave_If(self, original_node: cst.If) -> None: + self.if_else_depth -= 1 + + def visit_Else(self, node: cst.Else) -> Optional[bool]: + # Else blocks are already counted as part of the if statement + return True + + def visit_Assign(self, node: cst.Assign) -> Optional[bool]: + # Only process global assignments (not inside functions, classes, etc.) + if self.scope_depth == 0 and self.if_else_depth == 0: # We're at module level + for target in node.targets: + if isinstance(target.target, cst.Name): + name = target.target.value + self.assignments[name] = node + if name not in self.assignment_order: + self.assignment_order.append(name) + return True + + def visit_AnnAssign(self, node: cst.AnnAssign) -> Optional[bool]: + # Handle annotated assignments like: _CACHE: Dict[str, int] = {} + # Only process module-level annotated assignments with a value + if ( + self.scope_depth == 0 + and self.if_else_depth == 0 + and isinstance(node.target, cst.Name) + and node.value is not None + ): + name = node.target.value + self.assignments[name] = node + if name not in self.assignment_order: + self.assignment_order.append(name) + return True + + +def find_insertion_index_after_imports(node: cst.Module) -> int: + """Find the position of the last import statement in the top-level of the module.""" + insert_index = 0 + for i, stmt in enumerate(node.body): + is_top_level_import = isinstance(stmt, cst.SimpleStatementLine) and any( + isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body + ) + + is_conditional_import = isinstance(stmt, cst.If) and all( + isinstance(inner, cst.SimpleStatementLine) + and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) + for inner in stmt.body.body + ) + + if is_top_level_import or is_conditional_import: + insert_index = i + 1 + + # Stop scanning once we reach a class or function definition. + # Imports are supposed to be at the top of the file, but they can technically appear anywhere, even at the bottom of the file. + # Without this check, a stray import later in the file + # would incorrectly shift our insertion index below actual code definitions. + if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)): + break + + return insert_index + + +class GlobalAssignmentTransformer(cst.CSTTransformer): + """Transforms global assignments in the original file with those from the new file.""" + + def __init__(self, new_assignments: dict[str, cst.Assign | cst.AnnAssign], new_assignment_order: list[str]) -> None: + super().__init__() + self.new_assignments = new_assignments + self.new_assignment_order = new_assignment_order + self.processed_assignments: set[str] = set() + self.scope_depth = 0 + self.if_else_depth = 0 + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + self.scope_depth += 1 + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + self.scope_depth -= 1 + return updated_node + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + self.scope_depth += 1 + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: + self.scope_depth -= 1 + return updated_node + + def visit_If(self, node: cst.If) -> None: + self.if_else_depth += 1 + + def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If: + self.if_else_depth -= 1 + return updated_node + + def visit_Else(self, node: cst.Else) -> None: + # Else blocks are already counted as part of the if statement + pass + + def leave_Assign(self, original_node: cst.Assign, updated_node: cst.Assign) -> cst.CSTNode: + if self.scope_depth > 0 or self.if_else_depth > 0: + return updated_node + + # Check if this is a global assignment we need to replace + for target in original_node.targets: + if isinstance(target.target, cst.Name): + name = target.target.value + if name in self.new_assignments: + self.processed_assignments.add(name) + return self.new_assignments[name] + + return updated_node + + def leave_AnnAssign(self, original_node: cst.AnnAssign, updated_node: cst.AnnAssign) -> cst.CSTNode: + if self.scope_depth > 0 or self.if_else_depth > 0: + return updated_node + + # Check if this is a global annotated assignment we need to replace + if isinstance(original_node.target, cst.Name): + name = original_node.target.value + if name in self.new_assignments: + self.processed_assignments.add(name) + return self.new_assignments[name] + + return updated_node + + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + # Add any new assignments that weren't in the original file + new_statements = list(updated_node.body) + + # Find assignments to append + assignments_to_append = [ + (name, self.new_assignments[name]) + for name in self.new_assignment_order + if name not in self.processed_assignments and name in self.new_assignments + ] + + if not assignments_to_append: + return updated_node.with_changes(body=new_statements) + + # Collect all class and function names defined in the module + # These are the names that assignments might reference + module_defined_names: set[str] = set() + for stmt in new_statements: + if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)): + module_defined_names.add(stmt.name.value) + + # Partition assignments: those that reference module definitions go at the end, + # those that don't can go right after imports + assignments_after_imports: list[tuple[str, cst.Assign | cst.AnnAssign]] = [] + assignments_after_definitions: list[tuple[str, cst.Assign | cst.AnnAssign]] = [] + + for name, assignment in assignments_to_append: + # Get the value being assigned + if isinstance(assignment, (cst.Assign, cst.AnnAssign)) and assignment.value is not None: + value_node = assignment.value + else: + # No value to analyze, safe to place after imports + assignments_after_imports.append((name, assignment)) + continue + + # Collect names referenced in the assignment value + referenced_names = collect_referenced_names(value_node) + + # Check if any referenced names are module-level definitions + if referenced_names & module_defined_names: + # This assignment references a class/function, place it after definitions + assignments_after_definitions.append((name, assignment)) + else: + # Safe to place right after imports + assignments_after_imports.append((name, assignment)) + + # Insert assignments that don't depend on module definitions right after imports + if assignments_after_imports: + insert_index = find_insertion_index_after_imports(updated_node) + assignment_lines = [ + cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()]) + for _, assignment in assignments_after_imports + ] + new_statements = list(chain(new_statements[:insert_index], assignment_lines, new_statements[insert_index:])) + + # Insert assignments that depend on module definitions after all class/function definitions + if assignments_after_definitions: + # Find the position after the last function or class definition + insert_index = find_insertion_index_after_imports(cst.Module(body=new_statements)) + for i, stmt in enumerate(new_statements): + if isinstance(stmt, (cst.FunctionDef, cst.ClassDef)): + insert_index = i + 1 + + assignment_lines = [ + cst.SimpleStatementLine([assignment], leading_lines=[cst.EmptyLine()]) + for _, assignment in assignments_after_definitions + ] + new_statements = list(chain(new_statements[:insert_index], assignment_lines, new_statements[insert_index:])) + + return updated_node.with_changes(body=new_statements) + + +class GlobalStatementTransformer(cst.CSTTransformer): + """Transformer that appends global statements at the end of the module. + + This ensures that global statements (like function calls at module level) are placed + after all functions, classes, and assignments they might reference, preventing NameError + at module load time. + + This transformer should be run LAST after GlobalFunctionTransformer and + GlobalAssignmentTransformer have already added their content. + """ + + def __init__(self, global_statements: list[cst.SimpleStatementLine]) -> None: + super().__init__() + self.global_statements = global_statements + + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + if not self.global_statements: + return updated_node + + new_statements = list(updated_node.body) + + # Add empty line before each statement for readability + statement_lines = [ + stmt.with_changes(leading_lines=[cst.EmptyLine(), *stmt.leading_lines]) for stmt in self.global_statements + ] + + # Append statements at the end of the module + # This ensures they come after all functions, classes, and assignments + new_statements.extend(statement_lines) + + return updated_node.with_changes(body=new_statements) + + +class GlobalStatementCollector(cst.CSTVisitor): + """Visitor that collects all global statements (excluding imports and functions/classes).""" + + def __init__(self) -> None: + super().__init__() + self.global_statements = [] + self.in_function_or_class = False + + def visit_ClassDef(self, node: cst.ClassDef) -> bool: + # Don't visit inside classes + self.in_function_or_class = True + return False + + def leave_ClassDef(self, original_node: cst.ClassDef) -> None: + self.in_function_or_class = False + + def visit_FunctionDef(self, node: cst.FunctionDef) -> bool: + # Don't visit inside functions + self.in_function_or_class = True + return False + + def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: + self.in_function_or_class = False + + def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None: + if not self.in_function_or_class: + for statement in node.body: + # Skip imports and assignments (both regular and annotated) + if not isinstance(statement, (cst.Import, cst.ImportFrom, cst.Assign, cst.AnnAssign)): + self.global_statements.append(node) + break + + +class LastImportFinder(cst.CSTVisitor): + """Finds the position of the last import statement in the module.""" + + def __init__(self) -> None: + super().__init__() + self.last_import_line = 0 + self.current_line = 0 + + def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None: + self.current_line += 1 + for statement in node.body: + if isinstance(statement, (cst.Import, cst.ImportFrom)): + self.last_import_line = self.current_line + + +class DottedImportCollector(cst.CSTVisitor): + """Collects all top-level imports from a Python module in normalized dotted format, including top-level conditional imports like `if TYPE_CHECKING:`. + + Examples + -------- + import os ==> "os" + import dbt.adapters.factory ==> "dbt.adapters.factory" + from pathlib import Path ==> "pathlib.Path" + from recce.adapter.base import BaseAdapter ==> "recce.adapter.base.BaseAdapter" + from typing import Any, List, Optional ==> "typing.Any", "typing.List", "typing.Optional" + from recce.util.lineage import ( build_column_key, filter_dependency_maps) ==> "recce.util.lineage.build_column_key", "recce.util.lineage.filter_dependency_maps" + + """ + + def __init__(self) -> None: + self.imports: set[str] = set() + self.depth = 0 # top-level + + def get_full_dotted_name(self, expr: cst.BaseExpression) -> str: + if isinstance(expr, cst.Name): + return expr.value + if isinstance(expr, cst.Attribute): + return f"{self.get_full_dotted_name(expr.value)}.{expr.attr.value}" + return "" + + def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None: + for statement in block.body: + if isinstance(statement, cst.SimpleStatementLine): + for child in statement.body: + if isinstance(child, cst.Import): + for alias in child.names: + module = self.get_full_dotted_name(alias.name) + asname = alias.asname.name.value if alias.asname else alias.name.value + if isinstance(asname, cst.Attribute): + self.imports.add(module) + else: + self.imports.add(module if module == asname else f"{module}.{asname}") + + elif isinstance(child, cst.ImportFrom): + if child.module is None: + continue + module = self.get_full_dotted_name(child.module) + if isinstance(child.names, cst.ImportStar): + continue + for alias in child.names: + if isinstance(alias, cst.ImportAlias): + name = alias.name.value + asname = alias.asname.name.value if alias.asname else name + self.imports.add(f"{module}.{asname}") + + def visit_Module(self, node: cst.Module) -> None: + self.depth = 0 + self._collect_imports_from_block(node) + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + self.depth += 1 + + def leave_FunctionDef(self, node: cst.FunctionDef) -> None: + self.depth -= 1 + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + self.depth += 1 + + def leave_ClassDef(self, node: cst.ClassDef) -> None: + self.depth -= 1 + + def visit_If(self, node: cst.If) -> None: + if self.depth == 0: + self._collect_imports_from_block(node.body) + + def visit_Try(self, node: cst.Try) -> None: + if self.depth == 0: + self._collect_imports_from_block(node.body) + + +def extract_global_statements(source_code: str) -> tuple[cst.Module, list[cst.SimpleStatementLine]]: + """Extract global statements from source code.""" + module = cst.parse_module(source_code) + collector = GlobalStatementCollector() + module.visit(collector) + return module, collector.global_statements + + +def find_last_import_line(target_code: str) -> int: + """Find the line number of the last import statement.""" + module = cst.parse_module(target_code) + finder = LastImportFinder() + module.visit(finder) + return finder.last_import_line + + +class FutureAliasedImportTransformer(cst.CSTTransformer): + def leave_ImportFrom( + self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom + ) -> cst.BaseSmallStatement | cst.FlattenSentinel[cst.BaseSmallStatement] | cst.RemovalSentinel: + import libcst.matchers as m + + if ( + (updated_node_module := updated_node.module) + and updated_node_module.value == "__future__" + and all(m.matches(name, m.ImportAlias()) for name in updated_node.names) + ): + if names := [name for name in updated_node.names if name.asname is None]: + return updated_node.with_changes(names=names) + return cst.RemoveFromParent() + return updated_node + + +def delete___future___aliased_imports(module_code: str) -> str: + return cst.parse_module(module_code).visit(FutureAliasedImportTransformer()).code + + +def add_global_assignments(src_module_code: str, dst_module_code: str) -> str: + src_module, new_added_global_statements = extract_global_statements(src_module_code) + dst_module, existing_global_statements = extract_global_statements(dst_module_code) + + unique_global_statements = [] + for stmt in new_added_global_statements: + if any( + stmt is existing_stmt or stmt.deep_equals(existing_stmt) for existing_stmt in existing_global_statements + ): + continue + unique_global_statements.append(stmt) + + # Reuse already-parsed dst_module + original_module = dst_module + + # Parse the src_module_code once only (already done above: src_module) + # Collect assignments from the new file + new_assignment_collector = GlobalAssignmentCollector() + src_module.visit(new_assignment_collector) + + # Collect module-level functions from both source and destination + src_function_collector = GlobalFunctionCollector() + src_module.visit(src_function_collector) + + dst_function_collector = GlobalFunctionCollector() + original_module.visit(dst_function_collector) + + # Filter out functions that already exist in the destination (only add truly new functions) + new_functions = { + name: func + for name, func in src_function_collector.functions.items() + if name not in dst_function_collector.functions + } + new_function_order = [name for name in src_function_collector.function_order if name in new_functions] + + # If there are no assignments, no new functions, and no global statements, return unchanged + if not new_assignment_collector.assignments and not new_functions and not unique_global_statements: + return dst_module_code + + # The order of transformations matters: + # 1. Functions first - so assignments and statements can reference them + # 2. Assignments second - so they come after functions but before statements + # 3. Global statements last - so they can reference both functions and assignments + + # Transform functions if any + if new_functions: + function_transformer = GlobalFunctionTransformer(new_functions, new_function_order) + original_module = original_module.visit(function_transformer) + + # Transform assignments if any + if new_assignment_collector.assignments: + transformer = GlobalAssignmentTransformer( + new_assignment_collector.assignments, new_assignment_collector.assignment_order + ) + original_module = original_module.visit(transformer) + + # Insert global statements (like function calls at module level) LAST, + # after all functions and assignments are added, to ensure they can reference any + # functions or variables defined in the module + if unique_global_statements: + statement_transformer = GlobalStatementTransformer(unique_global_statements) + original_module = original_module.visit(statement_transformer) + + return original_module.code + + +def resolve_star_import(module_name: str, project_root: Path) -> set[str]: + try: + module_path = module_name.replace(".", "/") + possible_paths = [project_root / f"{module_path}.py", project_root / f"{module_path}/__init__.py"] + + module_file = None + for path in possible_paths: + if path.exists(): + module_file = path + break + + if module_file is None: + logger.warning(f"Could not find module file for {module_name}, skipping star import resolution") + return set() + + with module_file.open(encoding="utf8") as f: + module_code = f.read() + + tree = ast.parse(module_code) + + all_names = None + for node in ast.walk(tree): + if ( + isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == "__all__" + ): + if isinstance(node.value, (ast.List, ast.Tuple)): + all_names = [] + for elt in node.value.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + all_names.append(elt.value) + elif isinstance(elt, ast.Str): # Python < 3.8 compatibility + all_names.append(elt.s) + break + + if all_names is not None: + return set(all_names) + + public_names = set() + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + if not node.name.startswith("_"): + public_names.add(node.name) + elif isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and not target.id.startswith("_"): + public_names.add(target.id) + elif isinstance(node, ast.AnnAssign): + if isinstance(node.target, ast.Name) and not node.target.id.startswith("_"): + public_names.add(node.target.id) + elif isinstance(node, ast.Import) or ( + isinstance(node, ast.ImportFrom) and not any(alias.name == "*" for alias in node.names) + ): + for alias in node.names: + name = alias.asname or alias.name + if not name.startswith("_"): + public_names.add(name) + + return public_names + + except Exception as e: + logger.warning(f"Error resolving star import for {module_name}: {e}") + return set() + + +def add_needed_imports_from_module( + src_module_code: str, + dst_module_code: str | cst.Module, + src_path: Path, + dst_path: Path, + project_root: Path, + helper_functions: list[FunctionSource] | None = None, + helper_functions_fqn: set[str] | None = None, +) -> str: + """Add all needed and used source module code imports to the destination module code, and return it.""" + src_module_code = delete___future___aliased_imports(src_module_code) + if not helper_functions_fqn: + helper_functions_fqn = {f.fully_qualified_name for f in (helper_functions or [])} + + dst_code_fallback = dst_module_code if isinstance(dst_module_code, str) else dst_module_code.code + + src_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, src_path) + dst_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, dst_path) + + dst_context: CodemodContext = CodemodContext( + filename=src_path.name, + full_module_name=dst_module_and_package.name, + full_package_name=dst_module_and_package.package, + ) + gatherer: GatherImportsVisitor = GatherImportsVisitor( + CodemodContext( + filename=src_path.name, + full_module_name=src_module_and_package.name, + full_package_name=src_module_and_package.package, + ) + ) + try: + cst.parse_module(src_module_code).visit(gatherer) + except Exception as e: + logger.error(f"Error parsing source module code: {e}") + return dst_code_fallback + + dotted_import_collector = DottedImportCollector() + if isinstance(dst_module_code, cst.Module): + parsed_dst_module = dst_module_code + parsed_dst_module.visit(dotted_import_collector) + else: + try: + parsed_dst_module = cst.parse_module(dst_module_code) + parsed_dst_module.visit(dotted_import_collector) + except cst.ParserSyntaxError as e: + logger.exception(f"Syntax error in destination module code: {e}") + return dst_code_fallback + + try: + for mod in gatherer.module_imports: + # Skip __future__ imports as they cannot be imported directly + # __future__ imports should only be imported with specific objects i.e from __future__ import annotations + if mod == "__future__": + continue + if mod not in dotted_import_collector.imports: + AddImportsVisitor.add_needed_import(dst_context, mod) + RemoveImportsVisitor.remove_unused_import(dst_context, mod) + aliased_objects = set() + for mod, alias_pairs in gatherer.alias_mapping.items(): + for alias_pair in alias_pairs: + if alias_pair[0] and alias_pair[1]: # Both name and alias exist + aliased_objects.add(f"{mod}.{alias_pair[0]}") + + for mod, obj_seq in gatherer.object_mapping.items(): + for obj in obj_seq: + if ( + f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod # avoid circular deps + ): + continue # Skip adding imports for helper functions already in the context + + if f"{mod}.{obj}" in aliased_objects: + continue + + # Handle star imports by resolving them to actual symbol names + if obj == "*": + resolved_symbols = resolve_star_import(mod, project_root) + logger.debug(f"Resolved star import from {mod}: {resolved_symbols}") + + for symbol in resolved_symbols: + if ( + f"{mod}.{symbol}" not in helper_functions_fqn + and f"{mod}.{symbol}" not in dotted_import_collector.imports + ): + AddImportsVisitor.add_needed_import(dst_context, mod, symbol) + RemoveImportsVisitor.remove_unused_import(dst_context, mod, symbol) + else: + if f"{mod}.{obj}" not in dotted_import_collector.imports: + AddImportsVisitor.add_needed_import(dst_context, mod, obj) + RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj) + except Exception as e: + logger.exception(f"Error adding imports to destination module code: {e}") + return dst_code_fallback + + for mod, asname in gatherer.module_aliases.items(): + if not asname: + continue + if f"{mod}.{asname}" not in dotted_import_collector.imports: + AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname) + RemoveImportsVisitor.remove_unused_import(dst_context, mod, asname=asname) + + for mod, alias_pairs in gatherer.alias_mapping.items(): + for alias_pair in alias_pairs: + if f"{mod}.{alias_pair[0]}" in helper_functions_fqn: + continue + + if not alias_pair[0] or not alias_pair[1]: + continue + + if f"{mod}.{alias_pair[1]}" not in dotted_import_collector.imports: + AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1]) + RemoveImportsVisitor.remove_unused_import(dst_context, mod, alias_pair[0], asname=alias_pair[1]) + + try: + add_imports_visitor = AddImportsVisitor(dst_context) + transformed_module = add_imports_visitor.transform_module(parsed_dst_module) + transformed_module = RemoveImportsVisitor(dst_context).transform_module(transformed_module) + return transformed_module.code.lstrip("\n") + except Exception as e: + logger.exception(f"Error adding imports to destination module code: {e}") + return dst_code_fallback + + +def get_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | None, set[tuple[str, str]]]: + """Return the code for a function or methods in a Python module. + + functions_to_optimize is either a singleton FunctionToOptimize instance, which represents either a function at the + module level or a method of a class at the module level, or it represents a list of methods of the same class. + """ + if ( + not functions_to_optimize + or (functions_to_optimize[0].parents and functions_to_optimize[0].parents[0].type != "ClassDef") + or ( + len(functions_to_optimize[0].parents) > 1 + or ((len(functions_to_optimize) > 1) and len({fn.parents[0] for fn in functions_to_optimize}) != 1) + ) + ): + return None, set() + + file_path: Path = functions_to_optimize[0].file_path + class_skeleton: set[tuple[int, int | None]] = set() + contextual_dunder_methods: set[tuple[str, str]] = set() + target_code: str = "" + + def find_target(node_list: list[ast.stmt], name_parts: tuple[str, str] | tuple[str]) -> ast.AST | None: + target: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef | ast.Assign | ast.AnnAssign | None = None + node: ast.stmt + for node in node_list: + if ( + # The many mypy issues will be fixed once this code moves to the backend, + # using Type Guards as we move to 3.10+. + # We will cover the Type Alias case on the backend since it's a 3.12 feature. + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == name_parts[0] + ): + target = node + break + # The next two cases cover type aliases in pre-3.12 syntax, where only single assignment is allowed. + if ( + isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == name_parts[0] + ) or (isinstance(node, ast.AnnAssign) and hasattr(node.target, "id") and node.target.id == name_parts[0]): + if class_skeleton: + break + target = node + break + + if target is None or len(name_parts) == 1: + return target + + if not isinstance(target, ast.ClassDef) or len(name_parts) < 2: + return None + # At this point, name_parts has at least 2 elements + method_name: str = name_parts[1] # type: ignore[misc] + class_skeleton.add((target.lineno, target.body[0].lineno - 1)) + cbody = target.body + if isinstance(cbody[0], ast.expr): # Is a docstring + class_skeleton.add((cbody[0].lineno, cbody[0].end_lineno)) + cbody = cbody[1:] + cnode: ast.stmt + for cnode in cbody: + # Collect all dunder methods. + cnode_name: str + if ( + isinstance(cnode, (ast.FunctionDef, ast.AsyncFunctionDef)) + and len(cnode_name := cnode.name) > 4 + and cnode_name != method_name + and cnode_name.isascii() + and cnode_name.startswith("__") + and cnode_name.endswith("__") + ): + contextual_dunder_methods.add((target.name, cnode_name)) + class_skeleton.add((cnode.lineno, cnode.end_lineno)) + + return find_target(target.body, (method_name,)) + + with file_path.open(encoding="utf8") as file: + source_code: str = file.read() + try: + module_node: ast.Module = ast.parse(source_code) + except SyntaxError: + logger.exception("get_code - Syntax error while parsing code") + return None, set() + # Get the source code lines for the target node + lines: list[str] = source_code.splitlines(keepends=True) + if len(functions_to_optimize[0].parents) == 1: + if ( + functions_to_optimize[0].parents[0].type == "ClassDef" + ): # All functions_to_optimize functions are methods of the same class. + qualified_name_parts_list: list[tuple[str, str] | tuple[str]] = [ + (fto.parents[0].name, fto.function_name) for fto in functions_to_optimize + ] + + else: + logger.error(f"Error: get_code does not support inner functions: {functions_to_optimize[0].parents}") + return None, set() + elif len(functions_to_optimize[0].parents) == 0: + qualified_name_parts_list = [(functions_to_optimize[0].function_name,)] + else: + logger.error( + "Error: get_code does not support more than one level of nesting for now. " + f"Parents: {functions_to_optimize[0].parents}" + ) + return None, set() + for qualified_name_parts in qualified_name_parts_list: + target_node = find_target(module_node.body, qualified_name_parts) + if target_node is None: + continue + # find_target returns FunctionDef, AsyncFunctionDef, ClassDef, Assign, or AnnAssign - all have lineno/end_lineno + if not isinstance( + target_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Assign, ast.AnnAssign) + ): + continue + + if ( + isinstance(target_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) + and target_node.decorator_list + ): + target_code += "".join(lines[target_node.decorator_list[0].lineno - 1 : target_node.end_lineno]) + else: + target_code += "".join(lines[target_node.lineno - 1 : target_node.end_lineno]) + if not target_code: + return None, set() + class_list: list[tuple[int, int | None]] = sorted(class_skeleton) + class_code = "".join(["".join(lines[s_lineno - 1 : e_lineno]) for (s_lineno, e_lineno) in class_list]) + return class_code + target_code, contextual_dunder_methods + + +def extract_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | None, set[tuple[str, str]]]: + edited_code, contextual_dunder_methods = get_code(functions_to_optimize) + if edited_code is None: + return None, set() + try: + compile(edited_code, "edited_code", "exec") + except SyntaxError as e: + logger.exception(f"extract_code - Syntax error in extracted optimization candidate code: {e}") + return None, set() + return edited_code, contextual_dunder_methods + + +def find_preexisting_objects(source_code: str) -> set[tuple[str, tuple[FunctionParent, ...]]]: + """Find all preexisting functions, classes or class methods in the source code.""" + preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = set() + try: + module_node: ast.Module = ast.parse(source_code) + except SyntaxError: + logger.exception("find_preexisting_objects - Syntax error while parsing code") + return preexisting_objects + for node in module_node.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + preexisting_objects.add((node.name, ())) + elif isinstance(node, ast.ClassDef): + preexisting_objects.add((node.name, ())) + for cnode in node.body: + if isinstance(cnode, (ast.FunctionDef, ast.AsyncFunctionDef)): + preexisting_objects.add((cnode.name, (FunctionParent(node.name, "ClassDef"),))) + return preexisting_objects + + +@dataclass +class FunctionCallLocation: + """Represents a location where the target function is called.""" + + calling_function: str + line: int + column: int + + +@dataclass +class FunctionDefinitionInfo: + """Contains information about a function definition.""" + + name: str + node: ast.FunctionDef + source_code: str + start_line: int + end_line: int + is_method: bool + class_name: Optional[str] = None + + +class FunctionCallFinder(ast.NodeVisitor): + """AST visitor that finds all function definitions that call a specific qualified function. + + Args: + target_function_name: The qualified name of the function to find (e.g., "module.function" or "function") + target_filepath: The filepath where the target function is defined + + """ + + def __init__(self, target_function_name: str, target_filepath: str, source_lines: list[str]) -> None: + self.target_function_name = target_function_name + self.target_filepath = target_filepath + self.source_lines = source_lines # Store original source lines for extraction + + # Parse the target function name into parts + self.target_parts = target_function_name.split(".") + self.target_base_name = self.target_parts[-1] + + # Track current context + self.current_function_stack: list[tuple[str, ast.FunctionDef]] = [] + self.current_class_stack: list[str] = [] + + # Track imports to resolve qualified names + self.imports: dict[str, str] = {} # Maps imported names to their full paths + + # Results + self.function_calls: list[FunctionCallLocation] = [] + self.calling_functions: set[str] = set() + self.function_definitions: dict[str, FunctionDefinitionInfo] = {} + + # Track if we found calls in the current function + self.found_call_in_current_function = False + self.functions_with_nested_calls: set[str] = set() + + def visit_Import(self, node: ast.Import) -> None: + """Track regular imports.""" + for alias in node.names: + if alias.asname: + # import module as alias + self.imports[alias.asname] = alias.name + else: + # import module + self.imports[alias.name.split(".")[-1]] = alias.name + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + """Track from imports.""" + if node.module: + for alias in node.names: + if alias.name == "*": + # from module import * + self.imports["*"] = node.module + elif alias.asname: + # from module import name as alias + self.imports[alias.asname] = f"{node.module}.{alias.name}" + else: + # from module import name + self.imports[alias.name] = f"{node.module}.{alias.name}" + self.generic_visit(node) + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + """Track when entering a class definition.""" + self.current_class_stack.append(node.name) + self.generic_visit(node) + self.current_class_stack.pop() + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + """Track when entering a function definition.""" + self._visit_function_def(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + """Track when entering an async function definition.""" + self._visit_function_def(node) + + def _visit_function_def(self, node: ast.FunctionDef) -> None: + """Track when entering a function definition.""" + func_name = node.name + + # Build the full qualified name including class if applicable + full_name = f"{'.'.join(self.current_class_stack)}.{func_name}" if self.current_class_stack else func_name + + self.current_function_stack.append((full_name, node)) + self.found_call_in_current_function = False + + # Visit the function body + self.generic_visit(node) + + # Process the function after visiting its body + if self.found_call_in_current_function and full_name not in self.function_definitions: + # Extract function source code + source_code = self._extract_source_code(node) + + self.function_definitions[full_name] = FunctionDefinitionInfo( + name=full_name, + node=node, + source_code=source_code, + start_line=node.lineno, + end_line=node.end_lineno if hasattr(node, "end_lineno") else node.lineno, + is_method=bool(self.current_class_stack), + class_name=self.current_class_stack[-1] if self.current_class_stack else None, + ) + + # Handle nested functions - mark parent as containing nested calls + if self.found_call_in_current_function and len(self.current_function_stack) > 1: + parent_name = self.current_function_stack[-2][0] + self.functions_with_nested_calls.add(parent_name) + + # Also store the parent function if not already stored + if parent_name not in self.function_definitions: + parent_node = self.current_function_stack[-2][1] + parent_source = self._extract_source_code(parent_node) + + # Check if parent is a method (excluding current level) + parent_class_context = self.current_class_stack if len(self.current_function_stack) == 2 else [] + + self.function_definitions[parent_name] = FunctionDefinitionInfo( + name=parent_name, + node=parent_node, + source_code=parent_source, + start_line=parent_node.lineno, + end_line=parent_node.end_lineno if hasattr(parent_node, "end_lineno") else parent_node.lineno, + is_method=bool(parent_class_context), + class_name=parent_class_context[-1] if parent_class_context else None, + ) + + self.current_function_stack.pop() + + # Reset flag for parent function + if self.current_function_stack: + parent_name = self.current_function_stack[-1][0] + self.found_call_in_current_function = parent_name in self.calling_functions + + def visit_Call(self, node: ast.Call) -> None: + """Check if this call matches our target function.""" + if not self.current_function_stack: + # Not inside a function, skip + self.generic_visit(node) + return + + if self._is_target_function_call(node): + current_func_name = self.current_function_stack[-1][0] + + call_location = FunctionCallLocation( + calling_function=current_func_name, line=node.lineno, column=node.col_offset + ) + + self.function_calls.append(call_location) + self.calling_functions.add(current_func_name) + self.found_call_in_current_function = True + + self.generic_visit(node) + + def _is_target_function_call(self, node: ast.Call) -> bool: + """Determine if this call node is calling our target function.""" + call_name = self._get_call_name(node.func) + if not call_name: + return False + + # Check if it matches directly + if call_name == self.target_function_name: + return True + + # Check if it's just the base name matching + if call_name == self.target_base_name: + # Could be imported with a different name, check imports + if call_name in self.imports: + imported_path = self.imports[call_name] + if imported_path == self.target_function_name or imported_path.endswith( + f".{self.target_function_name}" + ): + return True + # Could also be a direct call if we're in the same file + return True + + # Check for qualified calls with imports + call_parts = call_name.split(".") + if call_parts[0] in self.imports: + # Resolve the full path using imports + base_import = self.imports[call_parts[0]] + full_path = f"{base_import}.{'.'.join(call_parts[1:])}" if len(call_parts) > 1 else base_import + + if full_path == self.target_function_name or full_path.endswith(f".{self.target_function_name}"): + return True + + return False + + def _get_call_name(self, func_node) -> Optional[str]: + """Extract the name being called from a function node.""" + # Fast path short-circuit for ast.Name nodes + if isinstance(func_node, ast.Name): + return func_node.id + + # Fast attribute chain extraction (speed: append, loop, join, NO reversed) + if isinstance(func_node, ast.Attribute): + parts = [] + current = func_node + # Unwind attribute chain as tight as possible (checked at each loop iteration) + while True: + parts.append(current.attr) + val = current.value + if isinstance(val, ast.Attribute): + current = val + continue + if isinstance(val, ast.Name): + parts.append(val.id) + # Join in-place backwards via slice instead of reversed for slight speedup + return ".".join(parts[::-1]) + break + return None + + def _extract_source_code(self, node: ast.FunctionDef) -> str: + """Extract source code for a function node using original source lines.""" + if not self.source_lines or not hasattr(node, "lineno"): + # Fallback to ast.unparse if available (Python 3.9+) + try: + return ast.unparse(node) + except AttributeError: + return f"# Source code extraction not available for {node.name}" + + # Get the lines for this function + start_line = node.lineno - 1 # Convert to 0-based index + end_line = node.end_lineno if hasattr(node, "end_lineno") else len(self.source_lines) + + # Extract the function lines + func_lines = self.source_lines[start_line:end_line] + + # Find the minimum indentation (excluding empty lines) + min_indent = float("inf") + for line in func_lines: + if line.strip(): # Skip empty lines + indent = len(line) - len(line.lstrip()) + min_indent = min(min_indent, indent) + + # If this is a method (inside a class), preserve one level of indentation + if self.current_class_stack: + # Keep 4 spaces of indentation for methods + dedent_amount = max(0, min_indent - 4) + result_lines = [] + for line in func_lines: + if line.strip(): # Only dedent non-empty lines + result_lines.append(line[dedent_amount:] if len(line) > dedent_amount else line) + else: + result_lines.append(line) + else: + # For top-level functions, remove all leading indentation + result_lines = [] + for line in func_lines: + if line.strip(): # Only dedent non-empty lines + result_lines.append(line[min_indent:] if len(line) > min_indent else line) + else: + result_lines.append(line) + + return "".join(result_lines).rstrip() + + def get_results(self) -> dict[str, str]: + """Get the results of the analysis. + + Returns: + A dictionary mapping qualified function names to their source code definitions. + + """ + return {info.name: info.source_code for info in self.function_definitions.values()} + + +def find_function_calls(source_code: str, target_function_name: str, target_filepath: str) -> dict[str, str]: + """Find all function definitions that call a specific target function. + + Args: + source_code: The Python source code to analyze + target_function_name: The qualified name of the function to find (e.g., "module.function") + target_filepath: The filepath where the target function is defined + + Returns: + A dictionary mapping qualified function names to their source code definitions. + Example: {"function_a": "def function_a(): ...", "MyClass.method_one": "def method_one(self): ..."} + + """ + # Parse the source code + tree = ast.parse(source_code) + + # Split source into lines for source extraction + source_lines = source_code.splitlines(keepends=True) + + # Create and run the visitor + visitor = FunctionCallFinder(target_function_name, target_filepath, source_lines) + visitor.visit(tree) + + return visitor.get_results() + + +def find_occurances( + qualified_name: str, file_path: str, fn_matches: list[Path], project_root: Path, tests_root: Path +) -> list[str]: # max chars for context + context_len = 0 + fn_call_context = "" + for cur_file in fn_matches: + if context_len > MAX_CONTEXT_LEN_REVIEW: + break + cur_file_path = Path(cur_file) + # exclude references in tests + try: + if cur_file_path.relative_to(tests_root): + continue + except ValueError: + pass + with cur_file_path.open(encoding="utf8") as f: + file_content = f.read() + results = find_function_calls(file_content, target_function_name=qualified_name, target_filepath=file_path) + if results: + try: + path_relative_to_project_root = cur_file_path.relative_to(project_root) + except Exception as e: + # shouldn't happen but ensuring we don't crash + logger.debug(f"investigate {e}") + continue + fn_call_context += f"```python:{path_relative_to_project_root}\n" + for ( + fn_definition + ) in results.values(): # multiple functions in the file might be calling the desired function + fn_call_context += f"{fn_definition}\n" + context_len += len(fn_definition) + fn_call_context += "```\n" + return fn_call_context + + +def find_specific_function_in_file( + source_code: str, filepath: Union[str, Path], target_function: str, target_class: str | None +) -> Optional[tuple[int, int]]: + """Find a specific function definition in a Python file and return its location. + + Stops searching once the target is found (optimized for performance). + + Args: + source_code: Source code string + filepath: Path to the Python file + target_function: Function Name of the function to find + target_class: Class name of the function to find + + Returns: + Tuple of (line_number, column_offset) if found, None otherwise + + """ + script = jedi.Script(code=source_code, path=filepath) + names = script.get_names(all_scopes=True, definitions=True) + for name in names: + if name.type == "function" and name.name == target_function: + # If class name specified, check parent + if target_class: + parent = name.parent() + if parent and parent.name == target_class and parent.type == "class": + return CodePosition(line_no=name.line, col_no=name.column) + else: + # Top-level function match + return CodePosition(line_no=name.line, col_no=name.column) + + return None # Function not found + + +def get_fn_references_jedi( + source_code: str, file_path: Path, project_root: Path, target_function: str, target_class: str | None +) -> list[Path]: + start_time = time.perf_counter() + function_position: CodePosition | None = find_specific_function_in_file( + source_code, file_path, target_function, target_class + ) + if function_position is None: + # Function not found (may be non-Python code) + return [] + try: + script = jedi.Script(code=source_code, path=file_path, project=jedi.Project(path=project_root)) + # Get references to the function + references = script.get_references(line=function_position.line_no, column=function_position.col_no) + # Collect unique file paths where references are found + end_time = time.perf_counter() + logger.debug(f"Jedi for function references ran in {end_time - start_time:.2f} seconds") + reference_files = set() + for ref in references: + if ref.module_path: + # Convert to string and normalize path + ref_path = str(ref.module_path) + # Skip the definition itself + if not (ref_path == file_path and ref.line == function_position.line_no): + reference_files.add(ref_path) + return sorted(reference_files) + except Exception as e: + print(f"Error during Jedi analysis: {e}") + return [] + + +has_numba = find_spec("numba") is not None + +NUMERICAL_MODULES = frozenset({"numpy", "torch", "numba", "jax", "tensorflow", "math", "scipy"}) +# Modules that require numba to be installed for optimization +NUMBA_REQUIRED_MODULES = frozenset({"numpy", "math", "scipy"}) + + +class NumericalUsageChecker(ast.NodeVisitor): + """AST visitor that checks if a function uses numerical computing libraries.""" + + def __init__(self, numerical_names: set[str]) -> None: + self.numerical_names = numerical_names + self.found_numerical = False + + def visit_Call(self, node: ast.Call) -> None: + """Check function calls for numerical library usage.""" + if self.found_numerical: + return + call_name = self._get_root_name(node.func) + if call_name and call_name in self.numerical_names: + self.found_numerical = True + return + self.generic_visit(node) + + def visit_Attribute(self, node: ast.Attribute) -> None: + """Check attribute access for numerical library usage.""" + if self.found_numerical: + return + root_name = self._get_root_name(node) + if root_name and root_name in self.numerical_names: + self.found_numerical = True + return + self.generic_visit(node) + + def visit_Name(self, node: ast.Name) -> None: + """Check name references for numerical library usage.""" + if self.found_numerical: + return + if node.id in self.numerical_names: + self.found_numerical = True + + def _get_root_name(self, node: ast.expr) -> str | None: + """Get the root name from an expression (e.g., 'np' from 'np.array').""" + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return self._get_root_name(node.value) + return None + + +def _collect_numerical_imports(tree: ast.Module) -> tuple[set[str], set[str]]: + """Collect names that reference numerical computing libraries from imports. + + Returns: + A tuple of (numerical_names, modules_used) where: + - numerical_names: set of names/aliases that reference numerical libraries + - modules_used: set of actual module names (e.g., "numpy", "math") being imported + + """ + numerical_names: set[str] = set() + modules_used: set[str] = set() + + stack: list[ast.AST] = [tree] + while stack: + node = stack.pop() + if isinstance(node, ast.Import): + for alias in node.names: + # import numpy or import numpy as np + module_root = alias.name.split(".")[0] + if module_root in NUMERICAL_MODULES: + # Use the alias if present, otherwise the module name + name = alias.asname if alias.asname else alias.name.split(".")[0] + numerical_names.add(name) + modules_used.add(module_root) + elif isinstance(node, ast.ImportFrom) and node.module: + module_root = node.module.split(".")[0] + if module_root in NUMERICAL_MODULES: + # from numpy import array, zeros as z + for alias in node.names: + if alias.name == "*": + # Can't track star imports, but mark the module as numerical + numerical_names.add(module_root) + else: + name = alias.asname if alias.asname else alias.name + numerical_names.add(name) + modules_used.add(module_root) + else: + stack.extend(ast.iter_child_nodes(node)) + + return numerical_names, modules_used + + +def _find_function_node(tree: ast.Module, name_parts: list[str]) -> ast.FunctionDef | None: + """Find a function node in the AST given its qualified name parts. + + Note: This function only finds regular (sync) functions, not async functions. + + Args: + tree: The parsed AST module + name_parts: List of name parts, e.g., ["ClassName", "method_name"] or ["function_name"] + + Returns: + The function node if found, None otherwise + + """ + if not name_parts: + return None + + if len(name_parts) == 1: + # Top-level function + func_name = name_parts[0] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == func_name: + return node + return None + + if len(name_parts) == 2: + # Class method: ClassName.method_name + class_name, method_name = name_parts + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for class_node in node.body: + if isinstance(class_node, ast.FunctionDef) and class_node.name == method_name: + return class_node + return None + + return None + + +def is_numerical_code(code_string: str, function_name: str | None = None) -> bool: + """Check if a function uses numerical computing libraries. + + Detects usage of numpy, torch, numba, jax, tensorflow, scipy, and math libraries + within the specified function. + + Note: For math, numpy, and scipy usage, this function returns True only if numba + is installed in the environment, as numba is required to optimize such code. + + Args: + code_string: The entire file's content as a string + function_name: The name of the function to check. Can be a simple name like "foo" + or a qualified name like "ClassName.method_name" for methods, + staticmethods, or classmethods. + + Returns: + True if the function uses any numerical computing library functions, False otherwise. + Returns False for math/numpy/scipy usage if numba is not installed. + + Examples: + >>> code = ''' + ... import numpy as np + ... def process_data(x): + ... return np.sum(x) + ... ''' + >>> is_numerical_code(code, "process_data") # Returns True only if numba is installed + True + + >>> code = ''' + ... def simple_func(x): + ... return x + 1 + ... ''' + >>> is_numerical_code(code, "simple_func") + False + + """ + try: + tree = ast.parse(code_string) + except SyntaxError: + return False + + # Collect names that reference numerical modules from imports + numerical_names, modules_used = _collect_numerical_imports(tree) + + if not function_name: + # Return True if modules used and (numba available or modules don't all require numba) + return bool(modules_used) and (has_numba or not modules_used.issubset(NUMBA_REQUIRED_MODULES)) + + # Split the function name to handle class methods + name_parts = function_name.split(".") + + # Find the target function node + target_function = _find_function_node(tree, name_parts) + if target_function is None: + return False + + # Check if the function body uses any numerical library + checker = NumericalUsageChecker(numerical_names) + checker.visit(target_function) + + if not checker.found_numerical: + return False + + # If numba is not installed and all modules used require numba for optimization, + # return False since we can't optimize this code + return not (not has_numba and modules_used.issubset(NUMBA_REQUIRED_MODULES)) + + +def get_opt_review_metrics( + source_code: str, file_path: Path, qualified_name: str, project_root: Path, tests_root: Path, language: Language +) -> str: + """Get function reference metrics for optimization review. + + Uses the LanguageSupport abstraction to find references, supporting both Python and JavaScript/TypeScript. + + Args: + source_code: Source code of the file containing the function. + file_path: Path to the file. + qualified_name: Qualified name of the function (e.g., "module.ClassName.method"). + project_root: Root of the project. + tests_root: Root of the tests directory. + language: The programming language. + + Returns: + Markdown-formatted string with code blocks showing calling functions. + + """ + from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.languages.registry import get_language_support + from codeflash.models.models import FunctionParent + + start_time = time.perf_counter() + + try: + # Get the language support + lang_support = get_language_support(language) + if lang_support is None: + return "" + + # Parse qualified name to get function name and class name + qualified_name_split = qualified_name.rsplit(".", maxsplit=1) + if len(qualified_name_split) == 1: + function_name, class_name = qualified_name_split[0], None + else: + function_name, class_name = qualified_name_split[1], qualified_name_split[0] + + # Create a FunctionToOptimize for the function + # We don't have full line info here, so we'll use defaults + parents: list[FunctionParent] = [] + if class_name: + parents = [FunctionParent(name=class_name, type="ClassDef")] + + func_info = FunctionToOptimize( + function_name=function_name, + file_path=file_path, + parents=parents, + starting_line=1, + ending_line=1, + language=str(language), + ) + + # Find references using language support + references = lang_support.find_references(func_info, project_root, tests_root, max_files=500) + + if not references: + return "" + + # Format references as markdown code blocks + calling_fns_details = _format_references_as_markdown(references, file_path, project_root, language) + + except Exception as e: + logger.debug(f"Error getting function references: {e}") + calling_fns_details = "" + + end_time = time.perf_counter() + logger.debug(f"Got function references in {end_time - start_time:.2f} seconds") + return calling_fns_details + + +def _format_references_as_markdown(references: list, file_path: Path, project_root: Path, language: Language) -> str: + """Format references as markdown code blocks with calling function code. + + Args: + references: List of ReferenceInfo objects. + file_path: Path to the source file (to exclude). + project_root: Root of the project. + language: The programming language. + + Returns: + Markdown-formatted string. + + """ + # Group references by file + refs_by_file: dict[Path, list] = {} + for ref in references: + # Exclude the source file's definition/import references + if ref.file_path == file_path and ref.reference_type in ("import", "reexport"): + continue + + if ref.file_path not in refs_by_file: + refs_by_file[ref.file_path] = [] + refs_by_file[ref.file_path].append(ref) + + from codeflash.languages.registry import get_language_support + + try: + lang_support = get_language_support(language) + except Exception: + lang_support = None + + fn_call_context = "" + context_len = 0 + + for ref_file, file_refs in refs_by_file.items(): + if context_len > MAX_CONTEXT_LEN_REVIEW: + break + + try: + path_relative = ref_file.relative_to(project_root) + except ValueError: + continue + + # Get syntax highlighting language + ext = ref_file.suffix.lstrip(".") + if language == Language.PYTHON: + lang_hint = "python" + elif ext in ("ts", "tsx"): + lang_hint = "typescript" + else: + lang_hint = "javascript" + + # Read the file to extract calling function context + try: + file_content = ref_file.read_text(encoding="utf-8") + lines = file_content.splitlines() + except Exception: + continue + + # Get unique caller functions from this file + callers_seen: set[str] = set() + caller_contexts: list[str] = [] + + for ref in file_refs: + caller = ref.caller_function or "" + if caller in callers_seen: + continue + callers_seen.add(caller) + + # Extract context around the reference + if ref.caller_function: + # Try to extract the full calling function + func_code = None + if lang_support is not None: + func_code = lang_support.extract_calling_function_source( + file_content, ref.caller_function, ref.line + ) + if func_code: + caller_contexts.append(func_code) + context_len += len(func_code) + else: + # Module-level call - show a few lines of context + start_line = max(0, ref.line - 3) + end_line = min(len(lines), ref.line + 2) + context_code = "\n".join(lines[start_line:end_line]) + caller_contexts.append(context_code) + context_len += len(context_code) + + if caller_contexts: + fn_call_context += f"```{lang_hint}:{path_relative.as_posix()}\n" + fn_call_context += "\n".join(caller_contexts) + fn_call_context += "\n```\n" + + return fn_call_context diff --git a/codeflash/languages/python/static_analysis/code_replacer.py b/codeflash/languages/python/static_analysis/code_replacer.py new file mode 100644 index 000000000..fd607d975 --- /dev/null +++ b/codeflash/languages/python/static_analysis/code_replacer.py @@ -0,0 +1,689 @@ +from __future__ import annotations + +import ast +from collections import defaultdict +from functools import lru_cache +from itertools import chain +from typing import TYPE_CHECKING, Optional, TypeVar + +import libcst as cst +from libcst.metadata import PositionProvider + +from codeflash.cli_cmds.console import logger +from codeflash.code_utils.config_parser import find_conftest_files +from codeflash.code_utils.formatter import sort_imports +from codeflash.languages import is_python +from codeflash.languages.python.static_analysis.code_extractor import ( + add_global_assignments, + add_needed_imports_from_module, + find_insertion_index_after_imports, +) +from codeflash.languages.python.static_analysis.line_profile_utils import ImportAdder +from codeflash.models.models import FunctionParent + +if TYPE_CHECKING: + from pathlib import Path + + from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.languages.base import LanguageSupport + from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, OptimizedCandidate, ValidCode + +ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST) + + +def normalize_node(node: ASTNodeT) -> ASTNodeT: + if isinstance(node, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and ast.get_docstring(node): + node.body = node.body[1:] + if hasattr(node, "body"): + node.body = [normalize_node(n) for n in node.body if not isinstance(n, (ast.Import, ast.ImportFrom))] + return node + + +@lru_cache(maxsize=3) +def normalize_code(code: str) -> str: + return ast.unparse(normalize_node(ast.parse(code))) + + +class AddRequestArgument(cst.CSTTransformer): + METADATA_DEPENDENCIES = (PositionProvider,) + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + # Matcher for '@fixture' or '@pytest.fixture' + for decorator in original_node.decorators: + dec = decorator.decorator + + if isinstance(dec, cst.Call): + func_name = "" + if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name): + if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest": + func_name = "pytest.fixture" + elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture": + func_name = "fixture" + + if func_name: + for arg in dec.args: + if ( + arg.keyword + and arg.keyword.value == "autouse" + and isinstance(arg.value, cst.Name) + and arg.value.value == "True" + ): + args = updated_node.params.params + arg_names = {arg.name.value for arg in args} + + # Skip if 'request' is already present + if "request" in arg_names: + return updated_node + + # Create a new 'request' param + request_param = cst.Param(name=cst.Name("request")) + + # Add 'request' as the first argument (after 'self' or 'cls' if needed) + if args: + first_arg = args[0].name.value + if first_arg in {"self", "cls"}: + new_params = [args[0], request_param] + list(args[1:]) # noqa: RUF005 + else: + new_params = [request_param] + list(args) # noqa: RUF005 + else: + new_params = [request_param] + + new_param_list = updated_node.params.with_changes(params=new_params) + return updated_node.with_changes(params=new_param_list) + return updated_node + + +class PytestMarkAdder(cst.CSTTransformer): + """Transformer that adds pytest marks to test functions.""" + + def __init__(self, mark_name: str) -> None: + super().__init__() + self.mark_name = mark_name + self.has_pytest_import = False + + def visit_Module(self, node: cst.Module) -> None: + """Check if pytest is already imported.""" + for statement in node.body: + if isinstance(statement, cst.SimpleStatementLine): + for stmt in statement.body: + if isinstance(stmt, cst.Import): + for import_alias in stmt.names: + if isinstance(import_alias, cst.ImportAlias) and import_alias.name.value == "pytest": + self.has_pytest_import = True + + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + """Add pytest import if not present.""" + if not self.has_pytest_import: + # Create import statement + import_stmt = cst.SimpleStatementLine(body=[cst.Import(names=[cst.ImportAlias(name=cst.Name("pytest"))])]) + # Add import at the beginning + updated_node = updated_node.with_changes(body=[import_stmt, *updated_node.body]) + return updated_node + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + """Add pytest mark to test functions.""" + # Check if the mark already exists + for decorator in updated_node.decorators: + if self._is_pytest_mark(decorator.decorator, self.mark_name): + return updated_node + + # Create the pytest mark decorator + mark_decorator = self._create_pytest_mark() + + # Add the decorator + new_decorators = [*list(updated_node.decorators), mark_decorator] + return updated_node.with_changes(decorators=new_decorators) + + def _is_pytest_mark(self, decorator: cst.BaseExpression, mark_name: str) -> bool: + """Check if a decorator is a specific pytest mark.""" + if isinstance(decorator, cst.Attribute): + if ( + isinstance(decorator.value, cst.Attribute) + and isinstance(decorator.value.value, cst.Name) + and decorator.value.value.value == "pytest" + and decorator.value.attr.value == "mark" + and decorator.attr.value == mark_name + ): + return True + elif isinstance(decorator, cst.Call) and isinstance(decorator.func, cst.Attribute): + return self._is_pytest_mark(decorator.func, mark_name) + return False + + def _create_pytest_mark(self) -> cst.Decorator: + """Create a pytest mark decorator.""" + # Base: pytest.mark.{mark_name} + mark_attr = cst.Attribute( + value=cst.Attribute(value=cst.Name("pytest"), attr=cst.Name("mark")), attr=cst.Name(self.mark_name) + ) + decorator = mark_attr + return cst.Decorator(decorator=decorator) + + +class AutouseFixtureModifier(cst.CSTTransformer): + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + # Matcher for '@fixture' or '@pytest.fixture' + for decorator in original_node.decorators: + dec = decorator.decorator + + if isinstance(dec, cst.Call): + func_name = "" + if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name): + if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest": + func_name = "pytest.fixture" + elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture": + func_name = "fixture" + + if func_name: + for arg in dec.args: + if ( + arg.keyword + and arg.keyword.value == "autouse" + and isinstance(arg.value, cst.Name) + and arg.value.value == "True" + ): + # Found a matching fixture with autouse=True + + # 1. The original body of the function will become the 'else' block. + # updated_node.body is an IndentedBlock, which is what cst.Else expects. + else_block = cst.Else(body=updated_node.body) + + # 2. Create the new 'if' block that will exit the fixture early. + if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")') + yield_statement = cst.parse_statement("yield") + if_body = cst.IndentedBlock(body=[yield_statement]) + + # 3. Construct the full if/else statement. + new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block) + + # 4. Replace the entire function's body with our new single statement. + return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement])) + return updated_node + + +def disable_autouse(test_path: Path) -> str: + file_content = test_path.read_text(encoding="utf-8") + module = cst.parse_module(file_content) + add_request_argument = AddRequestArgument() + disable_autouse_fixture = AutouseFixtureModifier() + modified_module = module.visit(add_request_argument) + modified_module = modified_module.visit(disable_autouse_fixture) + test_path.write_text(modified_module.code, encoding="utf-8") + return file_content + + +def modify_autouse_fixture(test_paths: list[Path]) -> dict[Path, list[str]]: + # find fixutre definition in conftetst.py (the one closest to the test) + # get fixtures present in override-fixtures in pyproject.toml + # add if marker closest return + file_content_map = {} + conftest_files = find_conftest_files(test_paths) + for cf_file in conftest_files: + # iterate over all functions in the file + # if function has autouse fixture, modify function to bypass with custom marker + original_content = disable_autouse(cf_file) + file_content_map[cf_file] = original_content + return file_content_map + + +# # reuse line profiler utils to add decorator and import to test fns +def add_custom_marker_to_all_tests(test_paths: list[Path]) -> None: + for test_path in test_paths: + # read file + file_content = test_path.read_text(encoding="utf-8") + module = cst.parse_module(file_content) + importadder = ImportAdder("import pytest") + modified_module = module.visit(importadder) + modified_module = cst.parse_module(sort_imports(code=modified_module.code, float_to_top=True)) + pytest_mark_adder = PytestMarkAdder("codeflash_no_autouse") + modified_module = modified_module.visit(pytest_mark_adder) + test_path.write_text(modified_module.code, encoding="utf-8") + + +def replace_functions_in_file( + source_code: str, + original_function_names: list[str], + optimized_code: str, + preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]], +) -> str: + parsed_function_names = [] + for original_function_name in original_function_names: + if original_function_name.count(".") == 0: + class_name, function_name = None, original_function_name + elif original_function_name.count(".") == 1: + class_name, function_name = original_function_name.split(".") + else: + msg = f"Unable to find {original_function_name}. Returning unchanged source code." + logger.error(msg) + return source_code + parsed_function_names.append((class_name, function_name)) + + # Collect functions from optimized code without using MetadataWrapper + optimized_module = cst.parse_module(optimized_code) + modified_functions: dict[tuple[str | None, str], cst.FunctionDef] = {} + new_functions: list[cst.FunctionDef] = [] + new_class_functions: dict[str, list[cst.FunctionDef]] = defaultdict(list) + new_classes: list[cst.ClassDef] = [] + modified_init_functions: dict[str, cst.FunctionDef] = {} + + function_names_set = set(parsed_function_names) + + for node in optimized_module.body: + if isinstance(node, cst.FunctionDef): + key = (None, node.name.value) + if key in function_names_set: + modified_functions[key] = node + elif preexisting_objects and (node.name.value, ()) not in preexisting_objects: + new_functions.append(node) + + elif isinstance(node, cst.ClassDef): + class_name = node.name.value + parents = (FunctionParent(name=class_name, type="ClassDef"),) + + if (class_name, ()) not in preexisting_objects: + new_classes.append(node) + + for child in node.body.body: + if isinstance(child, cst.FunctionDef): + method_key = (class_name, child.name.value) + if method_key in function_names_set: + modified_functions[method_key] = child + elif ( + child.name.value == "__init__" + and preexisting_objects + and (class_name, ()) in preexisting_objects + ): + modified_init_functions[class_name] = child + elif preexisting_objects and (child.name.value, parents) not in preexisting_objects: + new_class_functions[class_name].append(child) + + original_module = cst.parse_module(source_code) + + max_function_index = None + max_class_index = None + for index, _node in enumerate(original_module.body): + if isinstance(_node, cst.FunctionDef): + max_function_index = index + if isinstance(_node, cst.ClassDef): + max_class_index = index + + new_body: list[cst.CSTNode] = [] + existing_class_names = set() + + for node in original_module.body: + if isinstance(node, cst.FunctionDef): + key = (None, node.name.value) + if key in modified_functions: + modified_func = modified_functions[key] + new_body.append(node.with_changes(body=modified_func.body, decorators=modified_func.decorators)) + else: + new_body.append(node) + + elif isinstance(node, cst.ClassDef): + class_name = node.name.value + existing_class_names.add(class_name) + + new_members: list[cst.CSTNode] = [] + for child in node.body.body: + if isinstance(child, cst.FunctionDef): + key = (class_name, child.name.value) + if key in modified_functions: + modified_func = modified_functions[key] + new_members.append( + child.with_changes(body=modified_func.body, decorators=modified_func.decorators) + ) + elif child.name.value == "__init__" and class_name in modified_init_functions: + new_members.append(modified_init_functions[class_name]) + else: + new_members.append(child) + else: + new_members.append(child) + + if class_name in new_class_functions: + new_members.extend(new_class_functions[class_name]) + + new_body.append(node.with_changes(body=node.body.with_changes(body=new_members))) + else: + new_body.append(node) + + if new_classes: + unique_classes = [nc for nc in new_classes if nc.name.value not in existing_class_names] + if unique_classes: + new_classes_insertion_idx = ( + max_class_index if max_class_index is not None else find_insertion_index_after_imports(original_module) + ) + new_body = list( + chain(new_body[:new_classes_insertion_idx], unique_classes, new_body[new_classes_insertion_idx:]) + ) + + if new_functions: + if max_function_index is not None: + new_body = [*new_body[: max_function_index + 1], *new_functions, *new_body[max_function_index + 1 :]] + elif max_class_index is not None: + new_body = [*new_body[: max_class_index + 1], *new_functions, *new_body[max_class_index + 1 :]] + else: + new_body = [*new_functions, *new_body] + + updated_module = original_module.with_changes(body=new_body) + return updated_module.code + + +def replace_functions_and_add_imports( + source_code: str, + function_names: list[str], + optimized_code: str, + module_abspath: Path, + preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]], + project_root_path: Path, +) -> str: + return add_needed_imports_from_module( + optimized_code, + replace_functions_in_file(source_code, function_names, optimized_code, preexisting_objects), + module_abspath, + module_abspath, + project_root_path, + ) + + +def replace_function_definitions_in_module( + function_names: list[str], + optimized_code: CodeStringsMarkdown, + module_abspath: Path, + preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]], + project_root_path: Path, + should_add_global_assignments: bool = True, + function_to_optimize: Optional[FunctionToOptimize] = None, +) -> bool: + # Route to language-specific implementation for non-Python languages + if not is_python(): + return replace_function_definitions_for_language( + function_names, optimized_code, module_abspath, project_root_path, function_to_optimize + ) + + source_code: str = module_abspath.read_text(encoding="utf8") + code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code) + + new_code: str = replace_functions_and_add_imports( + # adding the global assignments before replacing the code, not after + # because of an "edge case" where the optimized code intoduced a new import and a global assignment using that import + # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet) + # this was added at https://github.com/codeflash-ai/codeflash/pull/448 + add_global_assignments(code_to_apply, source_code) if should_add_global_assignments else source_code, + function_names, + code_to_apply, + module_abspath, + preexisting_objects, + project_root_path, + ) + if is_zero_diff(source_code, new_code): + return False + module_abspath.write_text(new_code, encoding="utf8") + return True + + +def replace_function_definitions_for_language( + function_names: list[str], + optimized_code: CodeStringsMarkdown, + module_abspath: Path, + project_root_path: Path, + function_to_optimize: Optional[FunctionToOptimize] = None, +) -> bool: + """Replace function definitions for non-Python languages. + + Uses the language support abstraction to perform code replacement. + + Args: + function_names: List of qualified function names to replace. + optimized_code: The optimized code to apply. + module_abspath: Path to the module file. + project_root_path: Root of the project. + function_to_optimize: The function being optimized (needed for line info). + + Returns: + True if the code was modified, False if no changes. + + """ + from codeflash.languages import get_language_support + from codeflash.languages.base import Language + + original_source_code: str = module_abspath.read_text(encoding="utf8") + code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code) + + if not code_to_apply.strip(): + return False + + # Get language support + language = Language(optimized_code.language) + lang_support = get_language_support(language) + + # Add any new global declarations from the optimized code to the original source + original_source_code = lang_support.add_global_declarations( + optimized_code=code_to_apply, original_source=original_source_code, module_abspath=module_abspath + ) + + # If we have function_to_optimize with line info and this is the main file, use it for precise replacement + if ( + function_to_optimize + and function_to_optimize.starting_line + and function_to_optimize.ending_line + and function_to_optimize.file_path == module_abspath + ): + # Extract just the target function from the optimized code + optimized_func = _extract_function_from_code( + lang_support, code_to_apply, function_to_optimize.function_name, module_abspath + ) + if optimized_func: + new_code = lang_support.replace_function(original_source_code, function_to_optimize, optimized_func) + else: + # Fallback: use the entire optimized code (for simple single-function files) + new_code = lang_support.replace_function(original_source_code, function_to_optimize, code_to_apply) + else: + # For helper files or when we don't have precise line info: + # Find each function by name in both original and optimized code + # Then replace with the corresponding optimized version + new_code = original_source_code + modified = False + + # Get the list of function names to replace + functions_to_replace = list(function_names) + + for func_name in functions_to_replace: + # Re-discover functions from current code state to get correct line numbers + current_functions = lang_support.discover_functions_from_source(new_code, module_abspath) + + # Find the function in current code + func = None + for f in current_functions: + if func_name in (f.qualified_name, f.function_name): + func = f + break + + if func is None: + continue + + # Extract just this function from the optimized code + optimized_func = _extract_function_from_code( + lang_support, code_to_apply, func.function_name, module_abspath + ) + if optimized_func: + new_code = lang_support.replace_function(new_code, func, optimized_func) + modified = True + + if not modified: + logger.warning(f"Could not find function {function_names} in {module_abspath}") + return False + + # Check if there was actually a change + if original_source_code.strip() == new_code.strip(): + return False + + module_abspath.write_text(new_code, encoding="utf8") + return True + + +def _extract_function_from_code( + lang_support: LanguageSupport, source_code: str, function_name: str, file_path: Path | None = None +) -> str | None: + """Extract a specific function's source code from a code string. + + Includes JSDoc/docstring comments if present. + + Args: + lang_support: Language support instance. + source_code: The full source code containing the function. + function_name: Name of the function to extract. + file_path: Path to the file (used to determine correct analyzer for JS/TS). + + Returns: + The function's source code (including doc comments), or None if not found. + + """ + try: + # Use the language support to find functions in the source + # file_path is needed for JS/TS to determine correct analyzer (TypeScript vs JavaScript) + functions = lang_support.discover_functions_from_source(source_code, file_path) + for func in functions: + if func.function_name == function_name: + # Extract the function's source using line numbers + # Use doc_start_line if available to include JSDoc/docstring + lines = source_code.splitlines(keepends=True) + effective_start = func.doc_start_line or func.starting_line + if effective_start and func.ending_line and effective_start <= len(lines): + func_lines = lines[effective_start - 1 : func.ending_line] + return "".join(func_lines) + except Exception as e: + logger.debug(f"Error extracting function {function_name}: {e}") + + return None + + +def get_optimized_code_for_module(relative_path: Path, optimized_code: CodeStringsMarkdown) -> str: + file_to_code_context = optimized_code.file_to_path() + module_optimized_code = file_to_code_context.get(str(relative_path)) + if module_optimized_code is None: + # Fallback: if there's only one code block with None file path, + # use it regardless of the expected path (the AI server doesn't always include file paths) + if "None" in file_to_code_context and len(file_to_code_context) == 1: + module_optimized_code = file_to_code_context["None"] + logger.debug(f"Using code block with None file_path for {relative_path}") + else: + logger.warning( + f"Optimized code not found for {relative_path} In the context\n-------\n{optimized_code}\n-------\n" + "re-check your 'markdown code structure'" + f"existing files are {file_to_code_context.keys()}" + ) + module_optimized_code = "" + return module_optimized_code + + +def is_zero_diff(original_code: str, new_code: str) -> bool: + return normalize_code(original_code) == normalize_code(new_code) + + +def replace_optimized_code( + callee_module_paths: set[Path], + candidates: list[OptimizedCandidate], + code_context: CodeOptimizationContext, + function_to_optimize: FunctionToOptimize, + validated_original_code: dict[Path, ValidCode], + project_root: Path, +) -> tuple[set[Path], dict[str, dict[Path, str]]]: + initial_optimized_code = { + candidate.optimization_id: replace_functions_and_add_imports( + validated_original_code[function_to_optimize.file_path].source_code, + [function_to_optimize.qualified_name], + candidate.source_code, + function_to_optimize.file_path, + function_to_optimize.file_path, + code_context.preexisting_objects, + project_root, + ) + for candidate in candidates + } + callee_original_code = { + module_path: validated_original_code[module_path].source_code for module_path in callee_module_paths + } + intermediate_original_code: dict[str, dict[Path, str]] = { + candidate.optimization_id: ( + callee_original_code | {function_to_optimize.file_path: initial_optimized_code[candidate.optimization_id]} + ) + for candidate in candidates + } + module_paths = callee_module_paths | {function_to_optimize.file_path} + optimized_code = { + candidate.optimization_id: { + module_path: replace_functions_and_add_imports( + intermediate_original_code[candidate.optimization_id][module_path], + ( + [ + callee.qualified_name + for callee in code_context.helper_functions + if callee.file_path == module_path and callee.definition_type != "class" + ] + ), + candidate.source_code, + function_to_optimize.file_path, + module_path, + [], + project_root, + ) + for module_path in module_paths + } + for candidate in candidates + } + return module_paths, optimized_code + + +def is_optimized_module_code_zero_diff( + candidates: list[OptimizedCandidate], + validated_original_code: dict[Path, ValidCode], + optimized_code: dict[str, dict[Path, str]], + module_paths: set[Path], +) -> dict[str, dict[Path, bool]]: + return { + candidate.optimization_id: { + callee_module_path: normalize_code(optimized_code[candidate.optimization_id][callee_module_path]) + == validated_original_code[callee_module_path].normalized_code + for callee_module_path in module_paths + } + for candidate in candidates + } + + +def candidates_with_diffs( + candidates: list[OptimizedCandidate], + validated_original_code: ValidCode, + optimized_code: dict[str, dict[Path, str]], + module_paths: set[Path], +) -> list[OptimizedCandidate]: + return [ + candidate + for candidate in candidates + if not all( + is_optimized_module_code_zero_diff(candidates, validated_original_code, optimized_code, module_paths)[ + candidate.optimization_id + ].values() + ) + ] + + +def replace_optimized_code_in_worktrees( + optimized_code: dict[str, dict[Path, str]], + candidates: list[OptimizedCandidate], # Should be candidates_with_diffs + worktrees: list[Path], + git_root: Path, # Handle None case +) -> None: + for candidate, worktree in zip(candidates, worktrees[1:]): + for module_path in optimized_code[candidate.optimization_id]: + (worktree / module_path.relative_to(git_root)).write_text( + optimized_code[candidate.optimization_id][module_path], encoding="utf8" + ) # Check with is_optimized_module_code_zero_diff + + +def function_to_optimize_original_worktree_fqn( + function_to_optimize: FunctionToOptimize, worktrees: list[Path], git_root: Path +) -> str: + return ( + str(worktrees[0].name / function_to_optimize.file_path.relative_to(git_root).with_suffix("")).replace("/", ".") + + "." + + function_to_optimize.qualified_name + ) diff --git a/codeflash/languages/python/static_analysis/concolic_utils.py b/codeflash/languages/python/static_analysis/concolic_utils.py new file mode 100644 index 000000000..d674be370 --- /dev/null +++ b/codeflash/languages/python/static_analysis/concolic_utils.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import ast +import re +import subprocess +import uuid +from typing import Optional + +import sentry_sdk + +from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_temp_dir +from codeflash.code_utils.shell_utils import make_env_with_project_root + +# Known CrossHair limitations that produce invalid Python syntax in generated tests: +# - "" - higher-order functions returning nested functions +# - " object at 0x" - objects with default __repr__ +# - "", " object at 0x", " bool: + try: + ast.parse(test_code) + except SyntaxError: + is_known_limitation = any(pattern in test_code for pattern in CROSSHAIR_KNOWN_LIMITATION_PATTERNS) + if not is_known_limitation: + sentry_sdk.capture_message(f"CrossHair generated test with syntax error:\n{test_code}") + return False + + temp_path = (codeflash_temp_dir / f"concolic_test_{uuid.uuid4().hex}.py").resolve() + temp_path.write_text(test_code, encoding="utf-8") + + try: + result = subprocess.run( + [SAFE_SYS_EXECUTABLE, "-m", "pytest", "--collect-only", "-q", temp_path.as_posix()], + check=False, + capture_output=True, + text=True, + cwd=project_root, + timeout=10, + env=make_env_with_project_root(project_root) if project_root else None, + ) + except (subprocess.TimeoutExpired, Exception): + return False + else: + return result.returncode == 0 + finally: + temp_path.unlink(missing_ok=True) + + +class AssertCleanup: + def transform_asserts(self, code: str) -> str: + lines = code.splitlines() + result_lines = [] + + for line in lines: + transformed = self._transform_assert_line(line) + result_lines.append(transformed if transformed is not None else line) + + return "\n".join(result_lines) + + def _transform_assert_line(self, line: str) -> Optional[str]: + indent = line[: len(line) - len(line.lstrip())] + + assert_match = self.assert_re.match(line) + if assert_match: + expression = assert_match.group(1).strip() + if expression.startswith("not "): + return f"{indent}{expression}" + + expression = expression.rstrip(",;") + return f"{indent}{expression}" + + unittest_match = self.unittest_re.match(line) + if unittest_match: + indent, _assert_method, args = unittest_match.groups() + + if args: + arg_parts = self._first_top_level_arg(args) + if arg_parts: + return f"{indent}{arg_parts}" + + return None + + def __init__(self) -> None: + # Pre-compiling regular expressions for faster execution + self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$") + self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$") + + def _first_top_level_arg(self, args: str) -> str: + depth = 0 + for i, ch in enumerate(args): + if ch in "([{": + depth += 1 + elif ch in ")]}": + depth -= 1 + elif ch == "," and depth == 0: + return args[:i].strip() + return args.strip() + + +def clean_concolic_tests(test_suite_code: str) -> str: + try: + tree = ast.parse(test_suite_code) + can_parse = True + except Exception: + can_parse = False + tree = None + + if not can_parse or tree is None: + return AssertCleanup().transform_asserts(test_suite_code) + + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"): + new_body: list[ast.stmt] = [] + for stmt in node.body: + if isinstance(stmt, ast.Assert): + if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call): + new_body.append(ast.Expr(value=stmt.test.left)) + else: + new_body.append(stmt) + else: + new_body.append(stmt) + node.body = new_body + + return ast.unparse(tree).strip() diff --git a/codeflash/languages/python/static_analysis/coverage_utils.py b/codeflash/languages/python/static_analysis/coverage_utils.py new file mode 100644 index 000000000..b5d7ab8d8 --- /dev/null +++ b/codeflash/languages/python/static_analysis/coverage_utils.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import ast +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +from codeflash.code_utils.code_utils import get_run_tmp_file + +if TYPE_CHECKING: + from codeflash.models.models import CodeOptimizationContext + + +def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]: + """Extract the single dependent function from the code context excluding the main function.""" + dependent_functions = set() + + # Compare using bare name since AST extracts bare function names + bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function + + for code_string in code_context.testgen_context.code_strings: + # Quick heuristic: skip parsing entirely if there is no 'def' token, + # since no function definitions can be present without it. + if "def" not in code_string.code: + continue + + ast_tree = ast.parse(code_string.code) + # Add function names directly, skipping the bare main name. + for node in ast_tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + name = node.name + if name == bare_main: + continue + dependent_functions.add(name) + # If more than one dependent function (other than the main) is found, + # we can return False early since the final result cannot be a single name. + if len(dependent_functions) > 1: + return False + + if not dependent_functions: + return False + + if len(dependent_functions) != 1: + return False + + return build_fully_qualified_name(dependent_functions.pop(), code_context) + + +def build_fully_qualified_name(function_name: str, code_context: CodeOptimizationContext) -> str: + # If the name is already qualified (contains a dot), return as-is + if "." in function_name: + return function_name + full_name = function_name + for obj_name, parents in code_context.preexisting_objects: + if obj_name == function_name: + for parent in parents: + if parent.type == "ClassDef": + full_name = f"{parent.name}.{full_name}" + break + return full_name + + +def generate_candidates(source_code_path: Path) -> set[str]: + """Generate all the possible candidates for coverage data based on the source code path.""" + candidates = set() + # Add the filename as a candidate + name = source_code_path.name + candidates.add(name) + + # Precompute parts for efficient candidate path construction + parts = source_code_path.parts + n = len(parts) + + # Walk up the directory structure without creating Path objects or repeatedly converting to posix + last_added = name + # Start from the last parent and move up to the root, exclusive (skip the root itself) + for i in range(n - 2, 0, -1): + # Combine the ith part with the accumulated path (last_added) + candidate_path = f"{parts[i]}/{last_added}" + candidates.add(candidate_path) + last_added = candidate_path + + # Add the absolute posix path as a candidate + candidates.add(source_code_path.as_posix()) + return candidates + + +def prepare_coverage_files() -> tuple[Path, Path]: + """Prepare coverage configuration and output files.""" + coverage_database_file = get_run_tmp_file(Path(".coverage")) + coveragercfile = get_run_tmp_file(Path(".coveragerc")) + coveragerc_content = f"[run]\n branch = True\ndata_file={coverage_database_file}\n" + coveragercfile.write_text(coveragerc_content) + return coverage_database_file, coveragercfile diff --git a/codeflash/languages/python/static_analysis/edit_generated_tests.py b/codeflash/languages/python/static_analysis/edit_generated_tests.py new file mode 100644 index 000000000..c4aed07de --- /dev/null +++ b/codeflash/languages/python/static_analysis/edit_generated_tests.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +import ast +import os +import re +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +import libcst as cst +from libcst import MetadataWrapper +from libcst.metadata import PositionProvider + +from codeflash.cli_cmds.console import logger +from codeflash.code_utils.time_utils import format_perf, format_time +from codeflash.models.models import GeneratedTests, GeneratedTestsList +from codeflash.result.critic import performance_gain + +if TYPE_CHECKING: + from codeflash.models.models import InvocationId + + +class CommentMapper(ast.NodeVisitor): + def __init__( + self, test: GeneratedTests, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int] + ) -> None: + self.results: dict[int, str] = {} + self.test: GeneratedTests = test + self.original_runtimes = original_runtimes + self.optimized_runtimes = optimized_runtimes + self.abs_path = test.behavior_file_path.with_suffix("") + self.context_stack: list[str] = [] + + def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: + self.context_stack.append(node.name) + for inner_node in node.body: + if isinstance(inner_node, ast.FunctionDef): + self.visit_FunctionDef(inner_node) + elif isinstance(inner_node, ast.AsyncFunctionDef): + self.visit_AsyncFunctionDef(inner_node) + self.context_stack.pop() + return node + + def get_comment(self, match_key: str) -> str: + # calculate speedup and output comment + original_time = self.original_runtimes[match_key] + optimized_time = self.optimized_runtimes[match_key] + perf_gain = format_perf( + abs(performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) * 100) + ) + status = "slower" if optimized_time > original_time else "faster" + # Create the runtime comment + return f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})" + + def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: + self._process_function_def_common(node) + return node + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef: + self._process_function_def_common(node) + return node + + def _process_function_def_common(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + self.context_stack.append(node.name) + i = len(node.body) - 1 + test_qualified_name = ".".join(self.context_stack) + key = test_qualified_name + "#" + str(self.abs_path) + while i >= 0: + line_node = node.body[i] + if isinstance(line_node, (ast.With, ast.For, ast.While, ast.If)): + j = len(line_node.body) - 1 + while j >= 0: + compound_line_node: ast.stmt = line_node.body[j] + nodes_to_check = [compound_line_node] + nodes_to_check.extend(getattr(compound_line_node, "body", [])) + for internal_node in nodes_to_check: + if isinstance(internal_node, (ast.stmt, ast.Assign)): + inv_id = str(i) + "_" + str(j) + match_key = key + "#" + inv_id + if match_key in self.original_runtimes and match_key in self.optimized_runtimes: + self.results[internal_node.lineno] = self.get_comment(match_key) + j -= 1 + else: + inv_id = str(i) + match_key = key + "#" + inv_id + if match_key in self.original_runtimes and match_key in self.optimized_runtimes: + self.results[line_node.lineno] = self.get_comment(match_key) + i -= 1 + self.context_stack.pop() + + +def get_fn_call_linenos( + test: GeneratedTests, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int] +) -> dict[int, str]: + line_comment_ast_mapper = CommentMapper(test, original_runtimes, optimized_runtimes) + source_code = test.generated_original_test_source + tree = ast.parse(source_code) + line_comment_ast_mapper.visit(tree) + return line_comment_ast_mapper.results + + +class CommentAdder(cst.CSTTransformer): + """Transformer that adds comments to specified lines.""" + + # Declare metadata dependencies + METADATA_DEPENDENCIES = (PositionProvider,) + + def __init__(self, line_to_comments: dict[int, str]) -> None: + """Initialize the transformer with target line numbers. + + Args: + line_to_comments: Mapping of line numbers (1-indexed) to comments + + """ + self.line_to_comments = line_to_comments + super().__init__() + + def leave_SimpleStatementLine( + self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine + ) -> cst.SimpleStatementLine: + """Add comment to simple statement lines.""" + pos = self.get_metadata(PositionProvider, original_node) + + if pos and pos.start.line in self.line_to_comments: + # Create a comment with trailing whitespace + comment = cst.TrailingWhitespace( + whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment(self.line_to_comments[pos.start.line]) + ) + + # Update the trailing whitespace of the line itself + return updated_node.with_changes(trailing_whitespace=comment) + + return updated_node + + def leave_SimpleStatementSuite( + self, original_node: cst.SimpleStatementSuite, updated_node: cst.SimpleStatementSuite + ) -> cst.SimpleStatementSuite: + """Add comment to simple statement suites (e.g., after if/for/while).""" + pos = self.get_metadata(PositionProvider, original_node) + + if pos and pos.start.line in self.line_to_comments: + # Create a comment with trailing whitespace + comment = cst.TrailingWhitespace( + whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment(self.line_to_comments[pos.start.line]) + ) + + # Update the trailing whitespace of the suite + return updated_node.with_changes(trailing_whitespace=comment) + + return updated_node + + +def _is_python_file(file_path: Path) -> bool: + """Check if a file is a Python file.""" + return file_path.suffix == ".py" + + +def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path) -> dict[str, int]: + unique_inv_ids: dict[str, int] = {} + logger.debug(f"[unique_inv_id] Processing {len(inv_id_runtimes)} invocation IDs") + for inv_id, runtimes in inv_id_runtimes.items(): + test_qualified_name = ( + inv_id.test_class_name + "." + inv_id.test_function_name # type: ignore[operator] + if inv_id.test_class_name + else inv_id.test_function_name + ) + + test_module_path = inv_id.test_module_path + if "/" in test_module_path or "\\" in test_module_path: + abs_path = tests_project_rootdir / Path(test_module_path) + else: + abs_path = tests_project_rootdir / Path(test_module_path.replace(".", os.sep)).with_suffix(".py") + + abs_path_str = str(abs_path.resolve().with_suffix("")) + # Include both unit test and perf test paths for runtime annotations + # (performance test runtimes are used for annotations) + if ("__unit_test_" not in abs_path_str and "__perf_test_" not in abs_path_str) or not test_qualified_name: + logger.debug(f"[unique_inv_id] Skipping: path={abs_path_str}, test_qualified_name={test_qualified_name}") + continue + key = test_qualified_name + "#" + abs_path_str + parts = inv_id.iteration_id.split("_").__len__() # type: ignore[union-attr] + cur_invid = inv_id.iteration_id.split("_")[0] if parts < 3 else "_".join(inv_id.iteration_id.split("_")[:-1]) # type: ignore[union-attr] + match_key = key + "#" + cur_invid + logger.debug(f"[unique_inv_id] Adding key: {match_key} with runtime {min(runtimes)}") + if match_key not in unique_inv_ids: + unique_inv_ids[match_key] = 0 + unique_inv_ids[match_key] += min(runtimes) + logger.debug(f"[unique_inv_id] Result has {len(unique_inv_ids)} entries") + return unique_inv_ids + + +def add_runtime_comments_to_generated_tests( + generated_tests: GeneratedTestsList, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], + tests_project_rootdir: Optional[Path] = None, +) -> GeneratedTestsList: + """Add runtime performance comments to function calls in generated tests.""" + original_runtimes_dict = unique_inv_id(original_runtimes, tests_project_rootdir or Path()) + optimized_runtimes_dict = unique_inv_id(optimized_runtimes, tests_project_rootdir or Path()) + # Process each generated test + modified_tests = [] + for test in generated_tests.generated_tests: + is_python = _is_python_file(test.behavior_file_path) + + if is_python: + # Use Python libcst-based comment insertion + try: + tree = cst.parse_module(test.generated_original_test_source) + wrapper = MetadataWrapper(tree) + line_to_comments = get_fn_call_linenos(test, original_runtimes_dict, optimized_runtimes_dict) + comment_adder = CommentAdder(line_to_comments) + modified_tree = wrapper.visit(comment_adder) + modified_source = modified_tree.code + modified_test = GeneratedTests( + generated_original_test_source=modified_source, + instrumented_behavior_test_source=test.instrumented_behavior_test_source, + instrumented_perf_test_source=test.instrumented_perf_test_source, + behavior_file_path=test.behavior_file_path, + perf_file_path=test.perf_file_path, + ) + modified_tests.append(modified_test) + except Exception as e: + # If parsing fails, keep the original test + logger.debug(f"Failed to add runtime comments to test: {e}") + modified_tests.append(test) + else: + modified_tests.append(test) + + return GeneratedTestsList(generated_tests=modified_tests) + + +def remove_functions_from_generated_tests( + generated_tests: GeneratedTestsList, test_functions_to_remove: list[str] +) -> GeneratedTestsList: + # Pre-compile patterns for all function names to remove + function_patterns = _compile_function_patterns(test_functions_to_remove) + new_generated_tests = [] + + for generated_test in generated_tests.generated_tests: + source = generated_test.generated_original_test_source + + # Apply all patterns without redundant searches + for pattern in function_patterns: + # Use finditer and sub only if necessary to avoid unnecessary .search()/.sub() calls + for match in pattern.finditer(source): + # Skip if "@pytest.mark.parametrize" present + # Only the matched function's code is targeted + if "@pytest.mark.parametrize" in match.group(0): + continue + # Remove function from source + # If match, remove the function by substitution in the source + # Replace using start/end indices for efficiency + start, end = match.span() + source = source[:start] + source[end:] + # After removal, break since .finditer() is from left to right, and only one match expected per function in source + break + + generated_test.generated_original_test_source = source + new_generated_tests.append(generated_test) + + return GeneratedTestsList(generated_tests=new_generated_tests) + + +# Pre-compile all function removal regexes upfront for efficiency. +def _compile_function_patterns(test_functions_to_remove: list[str]) -> list[re.Pattern[str]]: + return [ + re.compile( + rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?(async\s+)?def\s+{re.escape(func)}\(.*?\):.*?(?=\n(async\s+)?def\s|$)", + re.DOTALL, + ) + for func in test_functions_to_remove + ] diff --git a/codeflash/languages/python/static_analysis/line_profile_utils.py b/codeflash/languages/python/static_analysis/line_profile_utils.py new file mode 100644 index 000000000..93997b2c6 --- /dev/null +++ b/codeflash/languages/python/static_analysis/line_profile_utils.py @@ -0,0 +1,388 @@ +"""Adapted from line_profiler (https://github.com/pyutils/line_profiler) written by Enthought, Inc. (BSD License).""" + +from __future__ import annotations + +import ast +from collections import defaultdict +from pathlib import Path +from typing import TYPE_CHECKING, Union + +import libcst as cst + +from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.code_utils.formatter import sort_imports + +if TYPE_CHECKING: + from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.models.models import CodeOptimizationContext + +# Known JIT decorators organized by module +# Format: {module_path: {decorator_name, ...}} +JIT_DECORATORS: dict[str, set[str]] = { + "numba": {"jit", "njit", "vectorize", "guvectorize", "stencil", "cfunc", "generated_jit"}, + "numba.cuda": {"jit"}, + "torch": {"compile"}, + "torch.jit": {"script", "trace"}, + "tensorflow": {"function"}, + "jax": {"jit"}, +} + + +class JitDecoratorDetector(ast.NodeVisitor): + """AST visitor that detects JIT compilation decorators considering import aliases.""" + + def __init__(self) -> None: + # Maps local name -> (module, original_name) + # e.g., {"nb": ("numba", None), "my_jit": ("numba", "jit")} + self.import_aliases: dict[str, tuple[str, str | None]] = {} + self.found_jit_decorator = False + + def visit_Import(self, node: ast.Import) -> None: + """Track regular imports like 'import numba' or 'import numba as nb'.""" + for alias in node.names: + # alias.name is the module name, alias.asname is the alias (or None) + local_name = alias.asname if alias.asname else alias.name + # For module imports, we store (module_name, None) to indicate it's a module import + self.import_aliases[local_name] = (alias.name, None) + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + """Track from imports like 'from numba import jit' or 'from numba import jit as my_jit'.""" + if node.module is None: + self.generic_visit(node) + return + + for alias in node.names: + local_name = alias.asname if alias.asname else alias.name + # For from imports, we store (module_name, imported_name) + self.import_aliases[local_name] = (node.module, alias.name) + self.generic_visit(node) + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + """Check function decorators for JIT decorators.""" + for decorator in node.decorator_list: + if self._is_jit_decorator(decorator): + self.found_jit_decorator = True + return + self.generic_visit(node) + + def _is_jit_decorator(self, node: ast.expr) -> bool: + """Check if a decorator node is a known JIT decorator.""" + # Handle Call nodes (e.g., @jit() or @numba.jit(nopython=True)) + if isinstance(node, ast.Call): + return self._is_jit_decorator(node.func) + + # Handle simple Name nodes (e.g., @jit when imported directly) + if isinstance(node, ast.Name): + return self._check_name_decorator(node.id) + + # Handle Attribute nodes (e.g., @numba.jit or @nb.jit) + if isinstance(node, ast.Attribute): + return self._check_attribute_decorator(node) + + return False + + def _check_name_decorator(self, name: str) -> bool: + """Check if a simple name decorator (e.g., @jit) is a JIT decorator.""" + if name not in self.import_aliases: + return False + + module, imported_name = self.import_aliases[name] + + if imported_name is None: + # This is a module import used as decorator (unlikely but possible) + return False + + # Check if this is a known JIT decorator from the module + return self._is_known_jit_decorator(module, imported_name) + + def _check_attribute_decorator(self, node: ast.Attribute) -> bool: + """Check if an attribute decorator (e.g., @numba.jit) is a JIT decorator.""" + # Build the full attribute chain + parts = self._get_attribute_parts(node) + if not parts: + return False + + # The first part might be an alias + first_part = parts[0] + rest_parts = parts[1:] + + # Check if first_part is an imported alias + if first_part in self.import_aliases: + module, imported_name = self.import_aliases[first_part] + + if imported_name is None: + # It's a module import (e.g., import numba as nb) + # The full path is module + rest_parts + if rest_parts: + full_module = module + decorator_name = rest_parts[-1] + if len(rest_parts) > 1: + full_module = f"{module}.{'.'.join(rest_parts[:-1])}" + return self._is_known_jit_decorator(full_module, decorator_name) + # It's a from import of something that has attributes + # e.g., from torch import jit; @jit.script + elif rest_parts: + full_module = f"{module}.{imported_name}" + decorator_name = rest_parts[-1] + if len(rest_parts) > 1: + full_module = f"{full_module}.{'.'.join(rest_parts[:-1])}" + return self._is_known_jit_decorator(full_module, decorator_name) + # first_part is used directly (e.g., @numba.jit without import alias) + # Reconstruct the full path + elif rest_parts: + full_module = first_part + if len(rest_parts) > 1: + full_module = f"{first_part}.{'.'.join(rest_parts[:-1])}" + decorator_name = rest_parts[-1] + return self._is_known_jit_decorator(full_module, decorator_name) + + return False + + def _get_attribute_parts(self, node: ast.Attribute) -> list[str]: + """Get all parts of an attribute chain (e.g., ['numba', 'cuda', 'jit']).""" + parts = [] + current = node + + while isinstance(current, ast.Attribute): + parts.append(current.attr) + current = current.value + + if isinstance(current, ast.Name): + parts.append(current.id) + parts.reverse() + return parts + + return [] + + def _is_known_jit_decorator(self, module: str, decorator_name: str) -> bool: + """Check if a decorator from a module is a known JIT decorator.""" + if module in JIT_DECORATORS: + return decorator_name in JIT_DECORATORS[module] + return False + + +def contains_jit_decorator(code: str) -> bool: + """Check if the code contains JIT compilation decorators from numba, torch, tensorflow, or jax. + + This function uses AST parsing to accurately detect JIT decorators even when: + - They are imported with aliases (e.g., import numba as nb; @nb.jit) + - They are imported directly (e.g., from numba import jit; @jit) + - They are called with arguments (e.g., @jit(nopython=True)) + """ + try: + tree = ast.parse(code) + except SyntaxError: + return False + + detector = JitDecoratorDetector() + detector.visit(tree) + return detector.found_jit_decorator + + +class LineProfilerDecoratorAdder(cst.CSTTransformer): + """Transformer that adds a decorator to a function with a specific qualified name.""" + + # TODO we don't support nested functions yet so they can only be inside classes, dont use qualified names, instead use the structure + def __init__(self, qualified_name: str, decorator_name: str) -> None: + """Initialize the transformer. + + Args: + ---- + qualified_name: The fully qualified name of the function to add the decorator to (e.g., "MyClass.nested_func.target_func"). + decorator_name: The name of the decorator to add. + + """ + super().__init__() + self.qualified_name_parts = qualified_name.split(".") + self.decorator_name = decorator_name + + # Track our current context path, only add when we encounter a class + self.context_stack = [] + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + # Track when we enter a class + self.context_stack.append(node.name.value) + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: + # Pop the context when we leave a class + self.context_stack.pop() + return updated_node + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + # Track when we enter a function + self.context_stack.append(node.name.value) + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: + # Check if the current context path matches our target qualified name + if self.context_stack == self.qualified_name_parts: + # Check if the decorator is already present + has_decorator = any( + self._is_target_decorator(decorator.decorator) for decorator in original_node.decorators + ) + + # Only add the decorator if it's not already there + if not has_decorator: + new_decorator = cst.Decorator(decorator=cst.Name(value=self.decorator_name)) + + # Add our new decorator to the existing decorators + updated_decorators = [new_decorator, *list(updated_node.decorators)] + updated_node = updated_node.with_changes(decorators=tuple(updated_decorators)) + + # Pop the context when we leave a function + self.context_stack.pop() + return updated_node + + def _is_target_decorator(self, decorator_node: Union[cst.Name, cst.Attribute, cst.Call]) -> bool: + """Check if a decorator matches our target decorator name.""" + if isinstance(decorator_node, cst.Name): + return decorator_node.value == self.decorator_name + if isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name): + return decorator_node.func.value == self.decorator_name + return False + + +class ProfileEnableTransformer(cst.CSTTransformer): + def __init__(self, filename: str) -> None: + # Flag to track if we found the import statement + self.found_import = False + # Track indentation of the import statement + self.import_indentation = None + self.filename = filename + + def leave_ImportFrom(self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom) -> cst.ImportFrom: + # Check if this is the line profiler import statement + if ( + isinstance(original_node.module, cst.Name) + and original_node.module.value == "line_profiler" + and any( + name.name.value == "profile" and (not name.asname or name.asname.name.value == "codeflash_line_profile") + for name in original_node.names + ) + ): + self.found_import = True + # Get the indentation from the original node + if hasattr(original_node, "leading_lines"): + leading_whitespace = original_node.leading_lines[-1].whitespace if original_node.leading_lines else "" + self.import_indentation = leading_whitespace + + return updated_node + + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + if not self.found_import: + return updated_node + + # Create a list of statements from the original module + new_body = list(updated_node.body) + + # Find the index of the import statement + import_index = None + for i, stmt in enumerate(new_body): + if isinstance(stmt, cst.SimpleStatementLine): + for small_stmt in stmt.body: + if isinstance(small_stmt, cst.ImportFrom) and ( + isinstance(small_stmt.module, cst.Name) + and small_stmt.module.value == "line_profiler" + and any( + name.name.value == "profile" + and (not name.asname or name.asname.name.value == "codeflash_line_profile") + for name in small_stmt.names + ) + ): + import_index = i + break + if import_index is not None: + break + + if import_index is not None: + # Create the new enable statement to insert after the import + enable_statement = cst.parse_statement(f"codeflash_line_profile.enable(output_prefix='{self.filename}')") + + # Insert the new statement after the import statement + new_body.insert(import_index + 1, enable_statement) + + # Create a new module with the updated body + return updated_node.with_changes(body=new_body) + + +def add_decorator_to_qualified_function(module: cst.Module, qualified_name: str, decorator_name: str) -> cst.Module: + """Add a decorator to a function with the exact qualified name in the source code. + + Args: + ---- + module: The Python source code as a CST module. + qualified_name: The fully qualified name of the function to add the decorator to (e.g., "MyClass.nested_func.target_func"). + decorator_name: The name of the decorator to add. + + Returns: + ------- + The modified CST module. + + """ + transformer = LineProfilerDecoratorAdder(qualified_name, decorator_name) + return module.visit(transformer) + + +def add_profile_enable(original_code: str, line_profile_output_file: str) -> str: + # TODO modify by using a libcst transformer + module = cst.parse_module(original_code) + transformer = ProfileEnableTransformer(line_profile_output_file) + modified_module = module.visit(transformer) + return modified_module.code + + +class ImportAdder(cst.CSTTransformer): + def __init__(self, import_statement) -> None: + self.import_statement = import_statement + self.has_import = False + + def leave_Module(self, original_node, updated_node): # noqa: ANN201 + # If the import is already there, don't add it again + if self.has_import: + return updated_node + + # Parse the import statement into a CST node + import_node = cst.parse_statement(self.import_statement) + + # Add the import to the module's body + return updated_node.with_changes(body=[import_node, *list(updated_node.body)]) + + def visit_ImportFrom(self, node) -> None: + # Check if the profile is already imported from line_profiler + if node.module and node.module.value == "line_profiler": + for import_alias in node.names: + if import_alias.name.value == "profile": + self.has_import = True + + +def add_decorator_imports(function_to_optimize: FunctionToOptimize, code_context: CodeOptimizationContext) -> Path: + """Add a profile decorator to a function in a Python file and all its helper functions.""" + # self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root + # grouped iteration, file to fns to optimize, from line_profiler import profile as codeflash_line_profile + file_paths = defaultdict(list) + line_profile_output_file = get_run_tmp_file(Path("baseline_lprof")) + file_paths[function_to_optimize.file_path].append(function_to_optimize.qualified_name) + for elem in code_context.helper_functions: + file_paths[elem.file_path].append(elem.qualified_name) + for file_path, fns_present in file_paths.items(): + # open file + file_contents = file_path.read_text("utf-8") + # parse to cst + module_node = cst.parse_module(file_contents) + for fn_name in fns_present: + # add decorator + module_node = add_decorator_to_qualified_function(module_node, fn_name, "codeflash_line_profile") + # add imports + # Create a transformer to add the import + transformer = ImportAdder("from line_profiler import profile as codeflash_line_profile") + # Apply the transformer to add the import + module_node = module_node.visit(transformer) + modified_code = sort_imports(code=module_node.code, float_to_top=True) + # write to file + with file_path.open("w", encoding="utf-8") as file: + file.write(modified_code) + # Adding profile.enable line for changing the savepath of the data, do this only for the main file and not the helper files + file_contents = function_to_optimize.file_path.read_text("utf-8") + modified_code = add_profile_enable(file_contents, line_profile_output_file.as_posix()) + function_to_optimize.file_path.write_text(modified_code, "utf-8") + return line_profile_output_file diff --git a/codeflash/languages/python/static_analysis/static_analysis.py b/codeflash/languages/python/static_analysis/static_analysis.py new file mode 100644 index 000000000..a0d04bfb1 --- /dev/null +++ b/codeflash/languages/python/static_analysis/static_analysis.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import ast +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +from pydantic import BaseModel, ConfigDict, field_validator + +if TYPE_CHECKING: + from codeflash.models.function_types import FunctionParent + + +ObjectDefT = TypeVar("ObjectDefT", ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef) + + +class ImportedInternalModuleAnalysis(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + full_name: str + file_path: Path + + @field_validator("name") + @classmethod + def name_is_identifier(cls, v: str) -> str: + if not v.isidentifier(): + msg = "must be an identifier" + raise ValueError(msg) + return v + + @field_validator("full_name") + @classmethod + def full_name_is_dotted_identifier(cls, v: str) -> str: + if any(not s or not s.isidentifier() for s in v.split(".")): + msg = "must be a dotted identifier" + raise ValueError(msg) + return v + + @field_validator("file_path") + @classmethod + def file_path_exists(cls, v: Path | None) -> Path | None: + if v and not v.exists(): + msg = "must be an existing path" + raise ValueError(msg) + return v + + +class FunctionKind(Enum): + FUNCTION = 0 + STATIC_METHOD = 1 + CLASS_METHOD = 2 + INSTANCE_METHOD = 3 + + +def parse_imports(code: str) -> list[ast.Import | ast.ImportFrom]: + return [node for node in ast.walk(ast.parse(code)) if isinstance(node, (ast.Import, ast.ImportFrom))] + + +def resolve_relative_name(module: str | None, level: int, current_module: str) -> str | None: + if level == 0: + return module + current_parts = current_module.split(".") + if level > len(current_parts): + return None + base_parts = current_parts[:-level] + if module: + base_parts.extend(module.split(".")) + return ".".join(base_parts) + + +def get_module_full_name(node: ast.Import | ast.ImportFrom, current_module: str) -> list[str]: + if isinstance(node, ast.Import): + return [alias.name for alias in node.names] + base_module = resolve_relative_name(node.module, node.level, current_module) + if base_module is None: + return [] + if node.module is None and node.level > 0: + return [f"{base_module}.{alias.name}" for alias in node.names] + return [base_module] + + +def is_internal_module(module_name: str, project_root: Path) -> bool: + module_path = module_name.replace(".", "/") + possible_paths = [project_root / f"{module_path}.py", project_root / module_path / "__init__.py"] + return any(path.exists() for path in possible_paths) + + +def get_module_file_path(module_name: str, project_root: Path) -> Path | None: + module_path = module_name.replace(".", "/") + possible_paths = [project_root / f"{module_path}.py", project_root / module_path / "__init__.py"] + for path in possible_paths: + if path.exists(): + return path.resolve() + return None + + +def analyze_imported_modules( + code_str: str, module_file_path: Path, project_root: Path +) -> list[ImportedInternalModuleAnalysis]: + """Statically finds and analyzes all imported internal modules.""" + module_rel_path = module_file_path.relative_to(project_root).with_suffix("") + current_module = ".".join(module_rel_path.parts) + imports = parse_imports(code_str) + module_names: set[str] = set() + for node in imports: + module_names.update(get_module_full_name(node, current_module)) + internal_modules = {module_name for module_name in module_names if is_internal_module(module_name, project_root)} + return [ + ImportedInternalModuleAnalysis(name=str(mod_name).split(".")[-1], full_name=mod_name, file_path=file_path) + for mod_name in internal_modules + if (file_path := get_module_file_path(mod_name, project_root)) is not None + ] + + +def get_first_top_level_object_def_ast( + object_name: str, object_type: type[ObjectDefT], node: ast.AST +) -> ObjectDefT | None: + for child in ast.iter_child_nodes(node): + if isinstance(child, object_type) and child.name == object_name: + return child + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + if descendant := get_first_top_level_object_def_ast(object_name, object_type, child): + return descendant + return None + + +def get_first_top_level_function_or_method_ast( + function_name: str, parents: list[FunctionParent], node: ast.AST +) -> ast.FunctionDef | ast.AsyncFunctionDef | None: + if not parents: + result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node) + if result is not None: + return result + return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, node) + if parents[0].type == "ClassDef" and ( + class_node := get_first_top_level_object_def_ast(parents[0].name, ast.ClassDef, node) + ): + result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node) + if result is not None: + return result + return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, class_node) + return None + + +def function_kind(node: ast.FunctionDef | ast.AsyncFunctionDef, parents: list[FunctionParent]) -> FunctionKind | None: + if not parents or parents[0].type in ["FunctionDef", "AsyncFunctionDef"]: + return FunctionKind.FUNCTION + if parents[0].type == "ClassDef": + for decorator in node.decorator_list: + if isinstance(decorator, ast.Name): + if decorator.id == "classmethod": + return FunctionKind.CLASS_METHOD + if decorator.id == "staticmethod": + return FunctionKind.STATIC_METHOD + return FunctionKind.INSTANCE_METHOD + return None + + +def has_typed_parameters(node: ast.FunctionDef | ast.AsyncFunctionDef, parents: list[FunctionParent]) -> bool: + kind = function_kind(node, parents) + if kind in [FunctionKind.FUNCTION, FunctionKind.STATIC_METHOD]: + return all(arg.annotation for arg in node.args.args) + if kind in [FunctionKind.CLASS_METHOD, FunctionKind.INSTANCE_METHOD]: + return all(arg.annotation for arg in node.args.args[1:]) + return False diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 58f66d0b8..cf55e6f61 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -21,9 +21,26 @@ if TYPE_CHECKING: from collections.abc import Sequence + from codeflash.languages.base import DependencyResolver + from codeflash.models.models import FunctionSource, GeneratedTestsList, InvocationId + logger = logging.getLogger(__name__) +def function_sources_to_helpers(sources: list[FunctionSource]) -> list[HelperFunction]: + return [ + HelperFunction( + name=fs.only_function_name, + qualified_name=fs.qualified_name, + file_path=fs.file_path, + source_code=fs.source_code, + start_line=fs.jedi_definition.line if fs.jedi_definition else 1, + end_line=fs.jedi_definition.line if fs.jedi_definition else 1, + ) + for fs in sources + ] + + @register_language class PythonSupport: """Python language support implementation. @@ -59,6 +76,37 @@ def test_framework(self) -> str: def comment_prefix(self) -> str: return "#" + @property + def dir_excludes(self) -> frozenset[str]: + return frozenset( + { + "__pycache__", + ".venv", + "venv", + ".tox", + ".nox", + ".eggs", + ".mypy_cache", + ".ruff_cache", + ".pytest_cache", + ".hypothesis", + "htmlcov", + ".pytype", + ".pyre", + ".pybuilder", + ".ipynb_checkpoints", + ".codeflash", + ".cache", + ".complexipy_cache", + "build", + "dist", + "sdist", + ".coverage*", + ".pyright*", + "*.egg-info", + } + ) + # === Discovery === def discover_functions( @@ -171,127 +219,39 @@ def discover_tests( # === Code Analysis === def extract_code_context(self, function: FunctionToOptimize, project_root: Path, module_root: Path) -> CodeContext: - """Extract function code and its dependencies. + """Extract function code and its dependencies via the canonical context pipeline.""" + from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context - Uses jedi and libcst for Python code analysis. - - Args: - function: The function to extract context for. - project_root: Root of the project. - module_root: Root of the module containing the function. - - Returns: - CodeContext with target code and dependencies. - - """ try: - source = function.file_path.read_text() + result = get_code_optimization_context(function, project_root) except Exception as e: - logger.exception("Failed to read %s: %s", function.file_path, e) + logger.warning("Failed to extract code context for %s: %s", function.function_name, e) return CodeContext(target_code="", target_file=function.file_path, language=Language.PYTHON) - # Extract the function source - lines = source.splitlines(keepends=True) - if function.starting_line and function.ending_line: - target_lines = lines[function.starting_line - 1 : function.ending_line] - target_code = "".join(target_lines) - else: - target_code = "" - - # Find helper functions - helpers = self.find_helper_functions(function, project_root) - - # Extract imports - import_lines = [] - for line in lines: - stripped = line.strip() - if stripped.startswith(("import ", "from ")): - import_lines.append(stripped) - elif stripped and not stripped.startswith("#"): - # Stop at first non-import, non-comment line - break + helpers = function_sources_to_helpers(result.helper_functions) return CodeContext( - target_code=target_code, + target_code=result.read_writable_code.markdown, target_file=function.file_path, helper_functions=helpers, - read_only_context="", - imports=import_lines, + read_only_context=result.read_only_context_code, + imports=[], language=Language.PYTHON, ) def find_helper_functions(self, function: FunctionToOptimize, project_root: Path) -> list[HelperFunction]: - """Find helper functions called by the target function. - - Uses jedi for Python code analysis. - - Args: - function: The target function to analyze. - project_root: Root of the project. - - Returns: - List of HelperFunction objects. - - """ - helpers: list[HelperFunction] = [] + """Find helper functions called by the target function via the canonical jedi pipeline.""" + from codeflash.languages.python.context.code_context_extractor import get_function_sources_from_jedi try: - import jedi - - from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages - from codeflash.optimization.function_context import belongs_to_function_qualified - - script = jedi.Script(path=function.file_path, project=jedi.Project(path=project_root)) - file_refs = script.get_names(all_scopes=True, definitions=False, references=True) - - qualified_name = function.qualified_name - - for ref in file_refs: - if not ref.full_name or not belongs_to_function_qualified(ref, qualified_name): - continue - - try: - definitions = ref.goto(follow_imports=True, follow_builtin_imports=False) - except Exception: - continue - - for definition in definitions: - definition_path = definition.module_path - if definition_path is None: - continue - - # Check if it's a valid helper (in project, not in target function) - is_valid = ( - str(definition_path).startswith(str(project_root)) - and not path_belongs_to_site_packages(definition_path) - and definition.full_name - and not belongs_to_function_qualified(definition, qualified_name) - and definition.type == "function" - ) - - if is_valid: - helper_qualified_name = get_qualified_name(definition.module_name, definition.full_name) - # Get source code - try: - helper_source = definition.get_line_code() - except Exception: - helper_source = "" - - helpers.append( - HelperFunction( - name=definition.name, - qualified_name=helper_qualified_name, - file_path=definition_path, - source_code=helper_source, - start_line=definition.line or 1, - end_line=definition.line or 1, - ) - ) - + _dict, sources = get_function_sources_from_jedi( + {function.file_path: {function.qualified_name}}, project_root + ) except Exception as e: logger.warning("Failed to find helpers for %s: %s", function.function_name, e) + return [] - return helpers + return function_sources_to_helpers(sources) def find_references( self, function: FunctionToOptimize, project_root: Path, tests_root: Path | None = None, max_files: int = 500 @@ -419,7 +379,7 @@ def replace_function(self, source: str, function: FunctionToOptimize, new_source Modified source code with function replaced. """ - from codeflash.code_utils.code_replacer import replace_functions_in_file + from codeflash.languages.python.static_analysis.code_replacer import replace_functions_in_file try: # Determine the function names to replace @@ -697,6 +657,59 @@ def leave_FunctionDef( except Exception: return test_source + def postprocess_generated_tests( + self, generated_tests: GeneratedTestsList, test_framework: str, project_root: Path, source_file_path: Path + ) -> GeneratedTestsList: + """Apply language-specific postprocessing to generated tests.""" + _ = test_framework, project_root, source_file_path + return generated_tests + + def remove_test_functions_from_generated_tests( + self, generated_tests: GeneratedTestsList, functions_to_remove: list[str] + ) -> GeneratedTestsList: + """Remove specific test functions from generated tests.""" + from codeflash.languages.python.static_analysis.edit_generated_tests import ( + remove_functions_from_generated_tests, + ) + + return remove_functions_from_generated_tests(generated_tests, functions_to_remove) + + def add_runtime_comments_to_generated_tests( + self, + generated_tests: GeneratedTestsList, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], + tests_project_rootdir: Path | None = None, + ) -> GeneratedTestsList: + """Add runtime comments to generated tests.""" + from codeflash.languages.python.static_analysis.edit_generated_tests import ( + add_runtime_comments_to_generated_tests, + ) + + return add_runtime_comments_to_generated_tests( + generated_tests, original_runtimes, optimized_runtimes, tests_project_rootdir + ) + + def add_global_declarations(self, optimized_code: str, original_source: str, module_abspath: Path) -> str: + _ = optimized_code, module_abspath + return original_source + + def extract_calling_function_source(self, source_code: str, function_name: str, ref_line: int) -> str | None: + """Extract the source code of a calling function in Python.""" + try: + import ast + + lines = source_code.splitlines() + tree = ast.parse(source_code) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_name: + end_line = node.end_lineno or node.lineno + if node.lineno <= ref_line <= end_line: + return "\n".join(lines[node.lineno - 1 : end_line]) + except Exception: + return None + return None + # === Test Result Comparison === def compare_test_results( @@ -728,15 +741,6 @@ def get_test_file_suffix(self) -> str: """ return ".py" - def get_comment_prefix(self) -> str: - """Get the comment prefix for Python. - - Returns: - Python single-line comment prefix. - - """ - return "#" - def find_test_root(self, project_root: Path) -> Path | None: """Find the test root directory for a Python project. @@ -801,6 +805,15 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: """ return True + def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | None: + from codeflash.languages.python.reference_graph import ReferenceGraph + + try: + return ReferenceGraph(project_root, language=self.language.value) + except Exception: + logger.debug("Failed to initialize ReferenceGraph, falling back to per-function Jedi analysis") + return None + def instrument_existing_test( self, test_path: Path,