wrapped functions default export support

Saga4 · Saga4 · commit b8597b2e85a2 · 2026-02-11T02:04:42.000+05:30
diff --git a/codeflash/languages/treesitter_utils.py b/codeflash/languages/treesitter_utils.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:
diff --git a/tests/test_languages/test_treesitter_utils.py b/tests/test_languages/test_treesitter_utils.py
@@ -693,3 +693,131 @@ def test_non_exported_const_not_exported(self, ts_analyzer):
         is_public_exported, name = ts_analyzer.is_function_exported(code, "publicFunc")
         assert is_public_exported is True
         assert name == "publicFunc"
+
+
+class TestWrappedDefaultExports:
+    """Tests for wrapped default export pattern - Issue #9.
+
+    Handles patterns like:
+    - export default curry(traverseEntity)
+    - export default compose(fn1, fn2)
+    - export default wrapper(myFunc)
+
+    These must be correctly recognized so the wrapped function is exportable.
+    """
+
+    @pytest.fixture
+    def ts_analyzer(self):
+        """Create a TypeScript analyzer."""
+        return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT)
+
+    def test_curry_wrapped_export(self, ts_analyzer):
+        """Test export default curry(fn) pattern."""
+        code = """import { curry } from 'lodash/fp';
+
+const traverseEntity = async (visitor, options, entity) => {
+    return entity;
+};
+
+export default curry(traverseEntity);"""
+
+        # Check exports parsing
+        exports = ts_analyzer.find_exports(code)
+        assert len(exports) == 1
+        assert exports[0].default_export == "default"
+        assert exports[0].wrapped_default_args == ["traverseEntity"]
+
+        # Check is_function_exported
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "traverseEntity")
+        assert is_exported is True
+        assert export_name == "default"
+
+    def test_compose_wrapped_export(self, ts_analyzer):
+        """Test export default compose(fn1, fn2) pattern with multiple args."""
+        code = """import { compose } from 'lodash/fp';
+
+function validateInput(data) { return data; }
+function processData(data) { return data; }
+
+export default compose(validateInput, processData);"""
+
+        exports = ts_analyzer.find_exports(code)
+        assert len(exports) == 1
+        assert exports[0].wrapped_default_args == ["validateInput", "processData"]
+
+        # Both functions should be recognized as exported
+        is_exported1, _ = ts_analyzer.is_function_exported(code, "validateInput")
+        is_exported2, _ = ts_analyzer.is_function_exported(code, "processData")
+        assert is_exported1 is True
+        assert is_exported2 is True
+
+    def test_nested_wrapper_export(self, ts_analyzer):
+        """Test nested wrapper: export default compose(curry(fn))."""
+        code = """export default compose(curry(myFunc));"""
+
+        exports = ts_analyzer.find_exports(code)
+        assert len(exports) == 1
+        assert "myFunc" in exports[0].wrapped_default_args
+
+        is_exported, _ = ts_analyzer.is_function_exported(code, "myFunc")
+        assert is_exported is True
+
+    def test_generic_wrapper_export(self, ts_analyzer):
+        """Test generic wrapper function."""
+        code = """const myFunction = (x: number) => x * 2;
+
+export default someWrapper(myFunction);"""
+
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "myFunction")
+        assert is_exported is True
+        assert export_name == "default"
+
+    def test_non_wrapped_function_not_exported(self, ts_analyzer):
+        """Test that functions not in the wrapper call are not exported."""
+        code = """const helper = (x: number) => x + 1;
+const main = (x: number) => helper(x) * 2;
+
+export default curry(main);"""
+
+        # main is wrapped, so it's exported
+        is_main_exported, _ = ts_analyzer.is_function_exported(code, "main")
+        assert is_main_exported is True
+
+        # helper is NOT in the wrapper call, so not exported
+        is_helper_exported, _ = ts_analyzer.is_function_exported(code, "helper")
+        assert is_helper_exported is False
+
+    def test_direct_default_export_still_works(self, ts_analyzer):
+        """Test that direct default exports still work."""
+        code = """function myFunc() { return 1; }
+export default myFunc;"""
+
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "myFunc")
+        assert is_exported is True
+        assert export_name == "default"
+
+    def test_strapi_traverse_entity_pattern(self, ts_analyzer):
+        """Test the exact strapi pattern that was failing."""
+        code = """import { curry } from 'lodash/fp';
+
+const traverseEntity = async (visitor: Visitor, options: TraverseOptions, entity: Data) => {
+    const { path = { raw: null }, schema, getModel } = options;
+    // ... implementation
+    return copy;
+};
+
+const createVisitorUtils = ({ data }: { data: Data }) => ({
+    remove(key: string) { delete data[key]; },
+    set(key: string, value: Data) { data[key] = value; },
+});
+
+export default curry(traverseEntity);"""
+
+        # traverseEntity should be recognized as exported
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "traverseEntity")
+        assert is_exported is True
+        assert export_name == "default"
+
+        # createVisitorUtils is NOT wrapped, so not exported via default
+        is_utils_exported, _ = ts_analyzer.is_function_exported(code, "createVisitorUtils")
+        assert is_utils_exported is False
