Fix arrow function and function expression export indexing

Arrow functions and function expressions assigned to variables
(e.g. `export const useAuth = () => { ... }`) were not being indexed
because the arrow_function AST node has no `name` field — the name
lives on the parent variable_declarator node.

Additionally, `isExported()` for TypeScript and JavaScript extractors
only checked 10 characters back from the node's start position, which
missed `export` for deeply nested nodes like arrow functions inside
variable declarations inside export statements.

Changes:
- extractFunction(): When an arrow_function or function_expression
  resolves to '<anonymous>', look up the parent variable_declarator
  for the name before skipping.
- isExported() (TS + JS): Walk the parent chain to find an
  export_statement ancestor instead of substring matching.
- Add 6 test cases covering arrow function exports, function
  expression exports, non-exported arrow functions, anonymous
  arrow functions, multiple exports, and JavaScript files.

Tested on a real monorepo (238 files): node count increased from
779 to 958 (+23%), with 94 new nodes in packages/ that previously
had 0 coverage.
This commit is contained in:
Tanner Balluff
2026-02-09 14:35:15 -05:00
parent f0ddfccf47
commit b6d5346c3c
2 changed files with 132 additions and 12 deletions
+97
View File
@@ -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 === '<anonymous>'
);
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 = `
+35 -12
View File
@@ -166,12 +166,17 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
}
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<Record<Language, LanguageExtractor>> = {
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 === '<anonymous>' &&
(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 === '<anonymous>') return; // Skip anonymous functions
const docstring = getPrecedingDocstring(node, this.source);