diff --git a/__tests__/extraction.test.ts b/__tests__/extraction.test.ts index ad0ba23..158cda8 100644 --- a/__tests__/extraction.test.ts +++ b/__tests__/extraction.test.ts @@ -745,6 +745,63 @@ export const fetchData = async () => { }); }); +describe('Generator Function Extraction (#1741)', () => { + const functionNames = (file: string, code: string) => + extractFromSource(file, code) + .nodes.filter((n) => n.kind === 'function') + .map((n) => n.name) + .sort(); + + it('extracts function* and async function* declarations in TypeScript', () => { + process.env.CODEGRAPH_KERNEL = '0'; + const code = ` +function plain() { return 1; } +function* gen() { yield 2; } +async function asyncFn() { return 3; } +async function* asyncGen() { yield 4; } +`; + expect(functionNames('gens.ts', code)).toEqual(['asyncFn', 'asyncGen', 'gen', 'plain']); + }); + + it('extracts function* and async function* declarations in JavaScript', () => { + process.env.CODEGRAPH_KERNEL = '0'; + const code = ` +function plain() { return 1; } +function* gen() { yield 2; } +async function asyncFn() { return 3; } +async function* asyncGen() { yield 4; } +`; + expect(functionNames('gens.js', code)).toEqual(['asyncFn', 'asyncGen', 'gen', 'plain']); + }); + + it('extracts const-assigned generator and async generator expressions (TS)', () => { + process.env.CODEGRAPH_KERNEL = '0'; + const code = ` +const g = function* () { yield 1; }; +const ag = async function* () { yield 2; }; +export const exportedGen = function* () { yield 3; }; +`; + const result = extractFromSource('gen-expr.ts', code); + const names = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name).sort(); + expect(names).toEqual(['ag', 'exportedGen', 'g']); + expect(result.nodes.find((n) => n.name === 'exportedGen')?.isExported).toBe(true); + expect(result.nodes.find((n) => n.name === 'g')?.isExported).toBeFalsy(); + }); + + it('extracts const-assigned generator and async generator expressions (JS)', () => { + process.env.CODEGRAPH_KERNEL = '0'; + const code = ` +const g = function* () { yield 1; }; +const ag = async function* () { yield 2; }; +export const exportedGen = function* () { yield 3; }; +`; + const result = extractFromSource('gen-expr.js', code); + const names = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name).sort(); + expect(names).toEqual(['ag', 'exportedGen', 'g']); + expect(result.nodes.find((n) => n.name === 'exportedGen')?.isExported).toBe(true); + }); +}); + describe('Type Alias Extraction', () => { it('should extract exported type aliases in TypeScript', () => { const code = ` diff --git a/codegraph-kernel/src/tsjs/extractors.rs b/codegraph-kernel/src/tsjs/extractors.rs index f240c5e..577a879 100644 --- a/codegraph-kernel/src/tsjs/extractors.rs +++ b/codegraph-kernel/src/tsjs/extractors.rs @@ -23,7 +23,7 @@ impl<'t> Walker<'t> { // variable_declarator (`export const useAuth = () => {}`). if name_override.is_none() && name == "" - && matches!(node.kind(), "arrow_function" | "function_expression") + && matches!(node.kind(), "arrow_function" | "function_expression" | "generator_function") { if let Some(parent) = node.parent() { if parent.kind() == "variable_declarator" { @@ -342,9 +342,9 @@ impl<'t> Walker<'t> { } let name = self.text(name_node).to_string(); - // Arrow/function values extract as functions, named by the declarator. + // Arrow/function/generator values extract as functions, named by the declarator. if let Some(v) = value { - if matches!(v.kind(), "arrow_function" | "function_expression") { + if matches!(v.kind(), "arrow_function" | "function_expression" | "generator_function") { self.extract_function(v, None); continue; } diff --git a/codegraph-kernel/src/tsjs/mod.rs b/codegraph-kernel/src/tsjs/mod.rs index afe6361..7ce4dcd 100644 --- a/codegraph-kernel/src/tsjs/mod.rs +++ b/codegraph-kernel/src/tsjs/mod.rs @@ -60,7 +60,7 @@ fn is_method_type(v: Variant, kind: &str) -> bool { } fn is_function_type(kind: &str) -> bool { - matches!(kind, "function_declaration" | "arrow_function" | "function_expression") + matches!(kind, "function_declaration" | "generator_function_declaration" | "arrow_function" | "function_expression" | "generator_function") } fn is_class_type(v: Variant, kind: &str) -> bool { @@ -792,7 +792,7 @@ impl<'t> Walker<'t> { if let Some(name_node) = node.child_by_field_name("name") { return self.text(name_node).to_string(); } - if matches!(node.kind(), "arrow_function" | "function_expression") { + if matches!(node.kind(), "arrow_function" | "function_expression" | "generator_function") { return "".to_string(); } for i in 0..node.named_child_count() { diff --git a/src/extraction/languages/javascript.ts b/src/extraction/languages/javascript.ts index 3b36348..899e153 100644 --- a/src/extraction/languages/javascript.ts +++ b/src/extraction/languages/javascript.ts @@ -3,7 +3,7 @@ import type { LanguageExtractor } from '../tree-sitter-types'; import { classifyTsClassMember } from './typescript'; export const javascriptExtractor: LanguageExtractor = { - functionTypes: ['function_declaration', 'arrow_function', 'function_expression'], + functionTypes: ['function_declaration', 'generator_function_declaration', 'arrow_function', 'function_expression', 'generator_function'], classTypes: ['class_declaration'], methodTypes: ['method_definition', 'field_definition'], // JS `field_definition` ≙ TS `public_field_definition`: plain fields are diff --git a/src/extraction/languages/typescript.ts b/src/extraction/languages/typescript.ts index 7205911..133aefd 100644 --- a/src/extraction/languages/typescript.ts +++ b/src/extraction/languages/typescript.ts @@ -39,7 +39,7 @@ export function classifyTsClassMember(node: SyntaxNode): 'method' | 'property' { } export const typescriptExtractor: LanguageExtractor = { - functionTypes: ['function_declaration', 'arrow_function', 'function_expression'], + functionTypes: ['function_declaration', 'generator_function_declaration', 'arrow_function', 'function_expression', 'generator_function'], classTypes: ['class_declaration', 'abstract_class_declaration'], methodTypes: ['method_definition', 'public_field_definition'], classifyMethodNode: classifyTsClassMember, diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index 7ef90c2..079b965 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -171,7 +171,7 @@ function extractNameRaw(node: SyntaxNode, source: string, extractor: LanguageExt // not from identifiers in their body. Without this, single-expression arrow // functions like `const fn = () => someIdentifier` get named "someIdentifier" // instead of "fn", because the fallback below finds the body identifier. - if (node.type === 'arrow_function' || node.type === 'function_expression') { + if (node.type === 'arrow_function' || node.type === 'function_expression' || node.type === 'generator_function') { return ''; } @@ -1556,7 +1556,7 @@ export class TreeSitterExtractor { if ( !nameOverride && name === '' && - (node.type === 'arrow_function' || node.type === 'function_expression') + (node.type === 'arrow_function' || node.type === 'function_expression' || node.type === 'generator_function') ) { const parent = node.parent; if (parent?.type === 'variable_declarator') { @@ -2623,7 +2623,7 @@ export class TreeSitterExtractor { } const name = getNodeText(nameNode, this.source); // Arrow functions / function expressions: extract as function instead of variable - if (valueNode && (valueNode.type === 'arrow_function' || valueNode.type === 'function_expression')) { + if (valueNode && (valueNode.type === 'arrow_function' || valueNode.type === 'function_expression' || valueNode.type === 'generator_function')) { this.extractFunction(valueNode); continue; }