Merge pull request #1441 from codeflash-ai/install_with_clone

Saga4 · web-flow · commit 97531dc62a35 · 2026-02-11T03:46:16.000+05:30
wrapped functions default export support
diff --git a/MULTI_LANGUAGE_ARCHITECTURE.md b/MULTI_LANGUAGE_ARCHITECTURE.md
@@ -386,7 +386,7 @@ class JavaScriptTransformer:
 
 from pathlib import Path
 from codeflash.languages.base import LanguageSupport, FunctionInfo, CodeContext
-from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
 from codeflash.languages.javascript.transformer import JavaScriptTransformer
 
 class JavaScriptSupport(LanguageSupport):
@@ -523,7 +523,7 @@ class JavaScriptSupport(LanguageSupport):
 # codeflash/languages/javascript/test_discovery.py
 
 from pathlib import Path
-from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
 
 class JestTestDiscovery:
     """Static analysis-based test discovery for Jest."""
diff --git a/codeflash/code_utils/code_extractor.py b/codeflash/code_utils/code_extractor.py
@@ -1772,7 +1772,7 @@ def _extract_calling_function_js(source_code: str, function_name: str, ref_line:
 
     """
     try:
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         # Try TypeScript first, fall back to JavaScript
         for lang in [TreeSitterLanguage.TYPESCRIPT, TreeSitterLanguage.TSX, TreeSitterLanguage.JAVASCRIPT]:
diff --git a/codeflash/code_utils/code_replacer.py b/codeflash/code_utils/code_replacer.py
@@ -26,7 +26,7 @@
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.languages.base import Language, LanguageSupport
-    from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+    from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
     from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, OptimizedCandidate, ValidCode
 
 ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)
@@ -640,7 +640,7 @@ def _add_global_declarations_for_language(
         return original_source
 
     try:
-        from codeflash.languages.treesitter_utils import get_analyzer_for_file
+        from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
         analyzer = get_analyzer_for_file(module_abspath)
 
diff --git a/codeflash/code_utils/normalizers/javascript.py b/codeflash/code_utils/normalizers/javascript.py
@@ -233,7 +233,7 @@ def normalize(self, code: str) -> str:
 
         """
         try:
-            from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+            from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
             lang_map = {"javascript": TreeSitterLanguage.JAVASCRIPT, "typescript": TreeSitterLanguage.TYPESCRIPT}
             lang = lang_map.get(self._get_tree_sitter_language(), TreeSitterLanguage.JAVASCRIPT)
diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py
@@ -201,7 +201,7 @@ def _is_js_ts_function_exported(file_path: Path, function_name: str) -> tuple[bo
         Tuple of (is_exported, export_name). export_name may be 'default' for default exports.
 
     """
-    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")
diff --git 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
+                    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",):
+        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
@@ -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__)
 
@@ -486,7 +486,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,8 +558,8 @@ def _find_helpers_recursive(
 
         """
         from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+        from codeflash.languages.javascript.treesitter import get_analyzer_for_file
         from codeflash.languages.registry import get_language_support
-        from codeflash.languages.treesitter_utils import get_analyzer_for_file
 
         if context.current_depth >= context.max_depth:
             return {}
diff --git a/codeflash/languages/javascript/instrument.py b/codeflash/languages/javascript/instrument.py
@@ -792,7 +792,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:
diff --git 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/support.py b/codeflash/languages/javascript/support.py
@@ -14,15 +14,15 @@
 
 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
 
 logger = logging.getLogger(__name__)
 
diff --git a/codeflash/languages/javascript/treesitter.py b/codeflash/languages/javascript/treesitter.py
@@ -94,6 +94,9 @@ class ExportInfo:
     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
@@ -707,6 +710,7 @@ def _extract_export_info(self, node: Node, source_bytes: bytes) -> ExportInfo |
         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")
@@ -726,6 +730,12 @@ def _extract_export_info(self, node: Node, source_bytes: bytes) -> ExportInfo |
                         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 }
@@ -773,8 +783,37 @@ def _extract_export_info(self, node: Node, source_bytes: bytes) -> ExportInfo |
             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.
 
@@ -876,6 +915,7 @@ def is_function_exported(
         """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.
@@ -901,6 +941,11 @@ def is_function_exported(
                 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:
@@ -1580,9 +1625,9 @@ def get_analyzer_for_file(file_path: Path) -> TreeSitterAnalyzer:
     """
     suffix = file_path.suffix.lower()
 
-    if suffix in (".ts",):
+    if suffix == ".ts":
         return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT)
-    if suffix in (".tsx",):
+    if suffix == ".tsx":
         return TreeSitterAnalyzer(TreeSitterLanguage.TSX)
     # Default to JavaScript for .js, .jsx, .mjs, .cjs
     return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)
diff --git a/codeflash/version.py b/codeflash/version.py
@@ -1,2 +1,2 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.20.0"
+__version__ = "0.20.0.post510.dev0+b8932209"
diff --git a/tests/test_languages/test_import_resolver.py b/tests/test_languages/test_import_resolver.py
@@ -8,7 +8,7 @@
 import pytest
 
 from codeflash.languages.javascript.import_resolver import HelperSearchContext, ImportResolver, MultiFileHelperFinder
-from codeflash.languages.treesitter_utils import ImportInfo
+from codeflash.languages.javascript.treesitter import ImportInfo
 
 
 class TestImportResolver:
@@ -286,7 +286,7 @@ class TestExportInfo:
     @pytest.fixture
     def js_analyzer(self):
         """Create a JavaScript analyzer."""
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)
 
@@ -388,7 +388,7 @@ class TestCommonJSRequire:
     @pytest.fixture
     def js_analyzer(self):
         """Create a JavaScript analyzer."""
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)
 
@@ -470,14 +470,14 @@ class TestCommonJSExports:
     @pytest.fixture
     def js_analyzer(self):
         """Create a JavaScript analyzer."""
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)
 
     @pytest.fixture
     def ts_analyzer(self):
         """Create a TypeScript analyzer."""
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT)
 
diff --git a/tests/test_languages/test_javascript_support.py b/tests/test_languages/test_javascript_support.py
@@ -654,7 +654,7 @@ def test_find_jest_tests(self, js_support):
             file_path = Path(f.name)
 
             source = file_path.read_text()
-            from codeflash.languages.treesitter_utils import get_analyzer_for_file
+            from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
             analyzer = get_analyzer_for_file(file_path)
             test_names = js_support._find_jest_tests(source, analyzer)
diff --git a/tests/test_languages/test_javascript_test_discovery.py b/tests/test_languages/test_javascript_test_discovery.py
diff --git a/tests/test_languages/test_treesitter_utils.py b/tests/test_languages/test_treesitter_utils.py
