Merge pull request #1246 from codeflash-ai/fix/class-method-export-detection

Fix class method export detection for JavaScript/TypeScript

Author: misrasaurabh1

codeflash/languages/javascript/find_references.py

@@ -101,6 +101,7 @@ def find_references(
         source_file: Path,
         include_definition: bool = False,
         max_files: int = 1000,
+        class_name: str | None = None,
     ) -> list[Reference]:
         """Find all references to a function across the project.
 
@@ -109,6 +110,7 @@ def find_references(
             source_file: Path to the file where the function is defined.
             include_definition: Whether to include the function definition itself.
             max_files: Maximum number of files to search (prevents runaway searches).
+            class_name: For class methods, the name of the containing class.
 
         Returns:
             List of Reference objects describing each call site.
@@ -126,7 +128,7 @@ def find_references(
             return references
 
         analyzer = get_analyzer_for_file(source_file)
-        exported = self._analyze_exports(function_name, source_file, source_code, analyzer)
+        exported = self._analyze_exports(function_name, source_file, source_code, analyzer, class_name)
 
         if not exported:
             logger.debug("Function %s is not exported from %s", function_name, source_file)
@@ -250,21 +252,29 @@ def find_references(
         return unique_refs
 
     def _analyze_exports(
-        self, function_name: str, file_path: Path, source_code: str, analyzer: TreeSitterAnalyzer
+        self,
+        function_name: str,
+        file_path: Path,
+        source_code: str,
+        analyzer: TreeSitterAnalyzer,
+        class_name: str | None = None,
     ) -> ExportedFunction | None:
         """Analyze how a function is exported from its file.
 
+        For class methods, also checks if the containing class is exported.
+
         Args:
             function_name: Name of the function to check.
             file_path: Path to the source file.
             source_code: Source code content.
             analyzer: TreeSitterAnalyzer instance.
+            class_name: For class methods, the name of the containing class.
 
         Returns:
             ExportedFunction if the function is exported, None otherwise.
 
         """
-        is_exported, export_name = analyzer.is_function_exported(source_code, function_name)
+        is_exported, export_name = analyzer.is_function_exported(source_code, function_name, class_name)
 
         if not is_exported:
             return None
@@ -825,6 +835,7 @@ def find_references(
     source_file: Path,
     project_root: Path | None = None,
     max_files: int = 1000,
+    class_name: str | None = None,
 ) -> list[Reference]:
     """Convenience function to find all references to a function.
 
@@ -835,6 +846,7 @@ def find_references(
         source_file: Path to the file where the function is defined.
         project_root: Root directory of the project. If None, uses source_file's parent.
         max_files: Maximum number of files to search.
+        class_name: For class methods, the name of the containing class.
 
     Returns:
         List of Reference objects describing each call site.
@@ -858,4 +870,4 @@ def find_references(
         project_root = source_file.parent
 
     finder = ReferenceFinder(project_root)
-    return finder.find_references(function_name, source_file, max_files=max_files)
+    return finder.find_references(function_name, source_file, max_files=max_files, class_name=class_name)

codeflash/languages/javascript/support.py

@@ -988,7 +988,9 @@ def find_references(
 
         try:
             finder = ReferenceFinder(project_root)
-            refs = finder.find_references(function.name, function.file_path, max_files=max_files)
+            refs = finder.find_references(
+                function.name, function.file_path, max_files=max_files, class_name=function.class_name
+            )
 
             # Convert to ReferenceInfo and filter out tests
             result: list[ReferenceInfo] = []

codeflash/languages/treesitter_utils.py

@@ -864,20 +864,27 @@ def _extract_commonjs_export(self, node: Node, source_bytes: bytes) -> ExportInf
             end_line=node.end_point[0] + 1,
         )
 
-    def is_function_exported(self, source: str, function_name: str) -> tuple[bool, str | None]:
+    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.
+
         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.
+            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:
@@ -888,6 +895,18 @@ def is_function_exported(self, source: str, function_name: str) -> tuple[bool, s
                 if name == function_name:
                     return (True, alias if alias else name)
 
+        # 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]:

tests/test_languages/test_import_resolver.py

@@ -474,6 +474,13 @@ def js_analyzer(self):
 
         return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)
 
+    @pytest.fixture
+    def ts_analyzer(self):
+        """Create a TypeScript analyzer."""
+        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+
+        return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT)
+
     def test_module_exports_function(self, js_analyzer):
         """Test module.exports = function() {}."""
         code = "module.exports = function helper() { return 1; };"
@@ -584,6 +591,51 @@ def test_is_function_exported_commonjs_property(self, js_analyzer):
         assert is_exported is True
         assert export_name == "helper"
 
+    def test_is_class_method_exported_via_class(self, ts_analyzer):
+        """Test is_function_exported returns True for method of exported class."""
+        code = """
+export class BloomFilter {
+    getHashValues(key: string): number[] {
+        return [1, 2, 3];
+    }
+}
+"""
+        # Method itself is not directly exported
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "getHashValues")
+        assert is_exported is False
+        assert export_name is None
+
+        # But when we pass the class name, it should find the class export
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "getHashValues", "BloomFilter")
+        assert is_exported is True
+        assert export_name == "BloomFilter"
+
+    def test_is_class_method_exported_default_class(self, ts_analyzer):
+        """Test is_function_exported returns True for method of default exported class."""
+        code = """
+export default class Calculator {
+    add(a: number, b: number): number {
+        return a + b;
+    }
+}
+"""
+        # When we pass the class name, it should find the default export
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "add", "Calculator")
+        assert is_exported is True
+        assert export_name == "Calculator"
+
+    def test_is_class_method_not_exported_non_exported_class(self, ts_analyzer):
+        """Test is_function_exported returns False for method of non-exported class."""
+        code = """
+class InternalClass {
+    helper(): void {}
+}
+"""
+        # Even with class name, non-exported class method should not be exported
+        is_exported, export_name = ts_analyzer.is_function_exported(code, "helper", "InternalClass")
+        assert is_exported is False
+        assert export_name is None
+
 
 class TestCommonJSImportResolver:
     """Tests for ImportResolver with CommonJS require() imports."""