diff --git a/__tests__/extraction.test.ts b/__tests__/extraction.test.ts index a1ac499..b88e395 100644 --- a/__tests__/extraction.test.ts +++ b/__tests__/extraction.test.ts @@ -193,6 +193,103 @@ function main() { }); }); +describe('Arrow Function Export Extraction', () => { + it('should extract exported arrow functions assigned to const', () => { + const code = ` +export const useAuth = (): AuthContextValue => { + return useContext(AuthContext); +}; +`; + const result = extractFromSource('hooks.ts', code); + + expect(result.nodes).toHaveLength(1); + expect(result.nodes[0]).toMatchObject({ + kind: 'function', + name: 'useAuth', + isExported: true, + }); + }); + + it('should extract exported function expressions assigned to const', () => { + const code = ` +export const processData = function(input: string): string { + return input.trim(); +}; +`; + const result = extractFromSource('utils.ts', code); + + expect(result.nodes).toHaveLength(1); + expect(result.nodes[0]).toMatchObject({ + kind: 'function', + name: 'processData', + isExported: true, + }); + }); + + it('should not extract non-exported arrow functions as exported', () => { + const code = ` +const internalHelper = () => { + return 42; +}; +`; + const result = extractFromSource('internal.ts', code); + + const helperNode = result.nodes.find((n) => n.name === 'internalHelper'); + expect(helperNode).toBeDefined(); + expect(helperNode?.isExported).toBeFalsy(); + }); + + it('should still skip truly anonymous arrow functions', () => { + const code = ` +const items = [1, 2, 3].map((x) => x * 2); +`; + const result = extractFromSource('anon.ts', code); + + // The inline arrow function passed to .map() has no variable_declarator parent + // and should remain anonymous (skipped) + const anonFunctions = result.nodes.filter( + (n) => n.kind === 'function' && n.name === '' + ); + expect(anonFunctions).toHaveLength(0); + }); + + it('should extract multiple exported arrow functions from the same file', () => { + const code = ` +export const add = (a: number, b: number): number => a + b; + +export const subtract = (a: number, b: number): number => a - b; + +const internal = () => 'not exported'; +`; + const result = extractFromSource('math.ts', code); + + const exported = result.nodes.filter((n) => n.kind === 'function' && n.isExported); + expect(exported).toHaveLength(2); + expect(exported.map((n) => n.name).sort()).toEqual(['add', 'subtract']); + + const internalNode = result.nodes.find((n) => n.name === 'internal'); + expect(internalNode).toBeDefined(); + expect(internalNode?.isExported).toBeFalsy(); + }); + + it('should extract arrow functions in JavaScript files', () => { + const code = ` +export const fetchData = async () => { + const response = await fetch('/api/data'); + return response.json(); +}; +`; + const result = extractFromSource('api.js', code); + + expect(result.nodes).toHaveLength(1); + expect(result.nodes[0]).toMatchObject({ + kind: 'function', + name: 'fetchData', + isExported: true, + }); + }); +}); + describe('Python Extraction', () => { it('should extract function definitions', () => { const code = ` diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index 7e91ecb..2dc26fe 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -166,12 +166,17 @@ const EXTRACTORS: Partial> = { } return undefined; }, - isExported: (node, source) => { - const parent = node.parent; - if (parent?.type === 'export_statement') return true; - // Check for 'export' keyword before declaration - const text = source.substring(Math.max(0, node.startIndex - 10), node.startIndex); - return text.includes('export'); + isExported: (node, _source) => { + // Walk the parent chain to find an export_statement ancestor. + // This correctly handles deeply nested nodes like arrow functions + // inside variable declarations: `export const X = () => { ... }` + // where the arrow_function is 3 levels deep under export_statement. + let current = node.parent; + while (current) { + if (current.type === 'export_statement') return true; + current = current.parent; + } + return false; }, isAsync: (node) => { for (let i = 0; i < node.childCount; i++) { @@ -204,11 +209,13 @@ const EXTRACTORS: Partial> = { const params = getChildByField(node, 'parameters'); return params ? getNodeText(params, source) : undefined; }, - isExported: (node, source) => { - const parent = node.parent; - if (parent?.type === 'export_statement') return true; - const text = source.substring(Math.max(0, node.startIndex - 10), node.startIndex); - return text.includes('export'); + isExported: (node, _source) => { + let current = node.parent; + while (current) { + if (current.type === 'export_statement') return true; + current = current.parent; + } + return false; }, isAsync: (node) => { for (let i = 0; i < node.childCount; i++) { @@ -891,7 +898,23 @@ export class TreeSitterExtractor { private extractFunction(node: SyntaxNode): void { if (!this.extractor) return; - const name = extractName(node, this.source, this.extractor); + let name = extractName(node, this.source, this.extractor); + // For arrow functions and function expressions assigned to variables, + // resolve the name from the parent variable_declarator. + // e.g. `export const useAuth = () => { ... }` — the arrow_function node + // has no `name` field; the name lives on the variable_declarator. + if ( + name === '' && + (node.type === 'arrow_function' || node.type === 'function_expression') + ) { + const parent = node.parent; + if (parent?.type === 'variable_declarator') { + const varName = getChildByField(parent, 'name'); + if (varName) { + name = getNodeText(varName, this.source); + } + } + } if (name === '') return; // Skip anonymous functions const docstring = getPrecedingDocstring(node, this.source);