feat(ipuaro): add decorator extraction to initial context

Extract decorators from classes and methods for NestJS/Angular support.
Decorators are now shown in initial context:
- @Controller('users') class UserController
- @Get(':id') async getUser(id: string): Promise<User>

Changes:
- Add decorators field to FunctionInfo, MethodInfo, ClassInfo
- Update ASTParser to extract decorators from tree-sitter nodes
- Update formatFileSummary to display decorators prefix
- Add 18 unit tests for decorator extraction and formatting
This commit is contained in:
imfozilbek
2025-12-05 13:38:46 +05:00
parent 5a22cd5c9b
commit 36768c06d1
6 changed files with 477 additions and 9 deletions

View File

@@ -52,6 +52,8 @@ export interface FunctionInfo {
isExported: boolean
/** Return type (if available) */
returnType?: string
/** Decorators applied to the function (e.g., ["@Get(':id')", "@Auth()"]) */
decorators?: string[]
}
export interface MethodInfo {
@@ -69,6 +71,8 @@ export interface MethodInfo {
visibility: "public" | "private" | "protected"
/** Whether it's static */
isStatic: boolean
/** Decorators applied to the method (e.g., ["@Get(':id')", "@UseGuards(AuthGuard)"]) */
decorators?: string[]
}
export interface PropertyInfo {
@@ -105,6 +109,8 @@ export interface ClassInfo {
isExported: boolean
/** Whether class is abstract */
isAbstract: boolean
/** Decorators applied to the class (e.g., ["@Controller('users')", "@Injectable()"]) */
decorators?: string[]
}
export interface InterfaceInfo {

View File

@@ -264,13 +264,15 @@ export class ASTParser {
const declaration = node.childForFieldName(FieldName.DECLARATION)
if (declaration) {
const decorators = this.extractDecoratorsFromSiblings(declaration)
switch (declaration.type) {
case NodeType.FUNCTION_DECLARATION:
this.extractFunction(declaration, ast, true)
this.extractFunction(declaration, ast, true, decorators)
this.addExportInfo(ast, declaration, "function", isDefault)
break
case NodeType.CLASS_DECLARATION:
this.extractClass(declaration, ast, true)
this.extractClass(declaration, ast, true, decorators)
this.addExportInfo(ast, declaration, "class", isDefault)
break
case NodeType.INTERFACE_DECLARATION:
@@ -309,7 +311,12 @@ export class ASTParser {
}
}
private extractFunction(node: SyntaxNode, ast: FileAST, isExported: boolean): void {
private extractFunction(
node: SyntaxNode,
ast: FileAST,
isExported: boolean,
externalDecorators: string[] = [],
): void {
const nameNode = node.childForFieldName(FieldName.NAME)
if (!nameNode) {
return
@@ -319,6 +326,9 @@ export class ASTParser {
const isAsync = node.children.some((c) => c.type === NodeType.ASYNC)
const returnTypeNode = node.childForFieldName(FieldName.RETURN_TYPE)
const nodeDecorators = this.extractNodeDecorators(node)
const decorators = [...externalDecorators, ...nodeDecorators]
ast.functions.push({
name: nameNode.text,
lineStart: node.startPosition.row + 1,
@@ -327,6 +337,7 @@ export class ASTParser {
isAsync,
isExported,
returnType: returnTypeNode?.text?.replace(/^:\s*/, ""),
decorators,
})
}
@@ -352,6 +363,7 @@ export class ASTParser {
isAsync,
isExported,
returnType: returnTypeNode?.text?.replace(/^:\s*/, ""),
decorators: [],
})
if (isExported) {
@@ -374,7 +386,12 @@ export class ASTParser {
}
}
private extractClass(node: SyntaxNode, ast: FileAST, isExported: boolean): void {
private extractClass(
node: SyntaxNode,
ast: FileAST,
isExported: boolean,
externalDecorators: string[] = [],
): void {
const nameNode = node.childForFieldName(FieldName.NAME)
if (!nameNode) {
return
@@ -385,14 +402,19 @@ export class ASTParser {
const properties: PropertyInfo[] = []
if (body) {
let pendingDecorators: string[] = []
for (const member of body.children) {
if (member.type === NodeType.METHOD_DEFINITION) {
methods.push(this.extractMethod(member))
if (member.type === NodeType.DECORATOR) {
pendingDecorators.push(this.formatDecorator(member))
} else if (member.type === NodeType.METHOD_DEFINITION) {
methods.push(this.extractMethod(member, pendingDecorators))
pendingDecorators = []
} else if (
member.type === NodeType.PUBLIC_FIELD_DEFINITION ||
member.type === NodeType.FIELD_DEFINITION
) {
properties.push(this.extractProperty(member))
pendingDecorators = []
}
}
}
@@ -400,6 +422,9 @@ export class ASTParser {
const { extendsName, implementsList } = this.extractClassHeritage(node)
const isAbstract = node.children.some((c) => c.type === NodeType.ABSTRACT)
const nodeDecorators = this.extractNodeDecorators(node)
const decorators = [...externalDecorators, ...nodeDecorators]
ast.classes.push({
name: nameNode.text,
lineStart: node.startPosition.row + 1,
@@ -410,6 +435,7 @@ export class ASTParser {
implements: implementsList,
isExported,
isAbstract,
decorators,
})
}
@@ -463,7 +489,7 @@ export class ASTParser {
}
}
private extractMethod(node: SyntaxNode): MethodInfo {
private extractMethod(node: SyntaxNode, decorators: string[] = []): MethodInfo {
const nameNode = node.childForFieldName(FieldName.NAME)
const params = this.extractParameters(node)
const isAsync = node.children.some((c) => c.type === NodeType.ASYNC)
@@ -485,6 +511,7 @@ export class ASTParser {
isAsync,
visibility,
isStatic,
decorators,
}
}
@@ -692,6 +719,49 @@ export class ASTParser {
}
}
/**
* Format a decorator node to a string like "@Get(':id')" or "@Injectable()".
*/
private formatDecorator(node: SyntaxNode): string {
return node.text.replace(/\s+/g, " ").trim()
}
/**
* Extract decorators that are direct children of a node.
* In tree-sitter, decorators are children of the class/function declaration.
*/
private extractNodeDecorators(node: SyntaxNode): string[] {
const decorators: string[] = []
for (const child of node.children) {
if (child.type === NodeType.DECORATOR) {
decorators.push(this.formatDecorator(child))
}
}
return decorators
}
/**
* Extract decorators from sibling nodes before the current node.
* Decorators appear as children before the declaration in export statements.
*/
private extractDecoratorsFromSiblings(node: SyntaxNode): string[] {
const decorators: string[] = []
const parent = node.parent
if (!parent) {
return decorators
}
for (const sibling of parent.children) {
if (sibling.type === NodeType.DECORATOR) {
decorators.push(this.formatDecorator(sibling))
} else if (sibling === node) {
break
}
}
return decorators
}
private classifyImport(from: string): ImportInfo["type"] {
if (from.startsWith(".") || from.startsWith("/")) {
return "internal"

View File

@@ -63,6 +63,9 @@ export const NodeType = {
DEFAULT: "default",
ACCESSIBILITY_MODIFIER: "accessibility_modifier",
READONLY: "readonly",
// Decorators
DECORATOR: "decorator",
} as const
export type NodeTypeValue = (typeof NodeType)[keyof typeof NodeType]

View File

@@ -187,10 +187,22 @@ function formatFileOverview(
return lines.join("\n")
}
/**
* Format decorators as a prefix string.
* Example: "@Get(':id') @Auth() "
*/
function formatDecoratorsPrefix(decorators: string[] | undefined): string {
if (!decorators || decorators.length === 0) {
return ""
}
return `${decorators.join(" ")} `
}
/**
* Format a function signature.
*/
function formatFunctionSignature(fn: FileAST["functions"][0]): string {
const decoratorsPrefix = formatDecoratorsPrefix(fn.decorators)
const asyncPrefix = fn.isAsync ? "async " : ""
const params = fn.params
.map((p) => {
@@ -200,7 +212,7 @@ function formatFunctionSignature(fn: FileAST["functions"][0]): string {
})
.join(", ")
const returnType = fn.returnType ? `: ${fn.returnType}` : ""
return `${asyncPrefix}${fn.name}(${params})${returnType}`
return `${decoratorsPrefix}${asyncPrefix}${fn.name}(${params})${returnType}`
}
/**
@@ -310,9 +322,10 @@ function formatFileSummary(
if (ast.classes.length > 0) {
for (const cls of ast.classes) {
const decoratorsPrefix = formatDecoratorsPrefix(cls.decorators)
const ext = cls.extends ? ` extends ${cls.extends}` : ""
const impl = cls.implements.length > 0 ? ` implements ${cls.implements.join(", ")}` : ""
lines.push(`- class ${cls.name}${ext}${impl}`)
lines.push(`- ${decoratorsPrefix}class ${cls.name}${ext}${impl}`)
}
}

View File

@@ -696,4 +696,140 @@ third: value3`
expect(ast.enums).toHaveLength(0)
})
})
describe("decorators (0.24.4)", () => {
it("should extract class decorator", () => {
const code = `@Controller('users')
class UserController {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].decorators).toHaveLength(1)
expect(ast.classes[0].decorators[0]).toBe("@Controller('users')")
})
it("should extract multiple class decorators", () => {
const code = `@Controller('api')
@Injectable()
@UseGuards(AuthGuard)
class ApiController {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].decorators).toHaveLength(3)
expect(ast.classes[0].decorators[0]).toBe("@Controller('api')")
expect(ast.classes[0].decorators[1]).toBe("@Injectable()")
expect(ast.classes[0].decorators[2]).toBe("@UseGuards(AuthGuard)")
})
it("should extract method decorators", () => {
const code = `class UserController {
@Get(':id')
@Auth()
async getUser() {}
}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].methods).toHaveLength(1)
expect(ast.classes[0].methods[0].decorators).toHaveLength(2)
expect(ast.classes[0].methods[0].decorators[0]).toBe("@Get(':id')")
expect(ast.classes[0].methods[0].decorators[1]).toBe("@Auth()")
})
it("should extract exported decorated class", () => {
const code = `@Injectable()
export class UserService {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].isExported).toBe(true)
expect(ast.classes[0].decorators).toHaveLength(1)
expect(ast.classes[0].decorators[0]).toBe("@Injectable()")
})
it("should extract decorator with complex arguments", () => {
const code = `@Module({
imports: [UserModule],
controllers: [AppController],
providers: [AppService]
})
class AppModule {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].decorators).toHaveLength(1)
expect(ast.classes[0].decorators[0]).toContain("@Module")
expect(ast.classes[0].decorators[0]).toContain("imports")
})
it("should extract decorated class with extends", () => {
const code = `@Entity()
class User extends BaseEntity {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].extends).toBe("BaseEntity")
expect(ast.classes[0].decorators).toHaveLength(1)
expect(ast.classes[0].decorators![0]).toBe("@Entity()")
})
it("should handle class without decorators", () => {
const code = `class SimpleClass {}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].decorators).toHaveLength(0)
})
it("should handle method without decorators", () => {
const code = `class SimpleClass {
simpleMethod() {}
}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].methods).toHaveLength(1)
expect(ast.classes[0].methods[0].decorators).toHaveLength(0)
})
it("should handle function without decorators", () => {
const code = `function simpleFunc() {}`
const ast = parser.parse(code, "ts")
expect(ast.functions).toHaveLength(1)
expect(ast.functions[0].decorators).toHaveLength(0)
})
it("should handle arrow function without decorators", () => {
const code = `const arrowFn = () => {}`
const ast = parser.parse(code, "ts")
expect(ast.functions).toHaveLength(1)
expect(ast.functions[0].decorators).toHaveLength(0)
})
it("should extract NestJS controller pattern", () => {
const code = `@Controller('users')
export class UserController {
@Get()
findAll() {}
@Get(':id')
findOne() {}
@Post()
@Body()
create() {}
}`
const ast = parser.parse(code, "ts")
expect(ast.classes).toHaveLength(1)
expect(ast.classes[0].decorators).toContain("@Controller('users')")
expect(ast.classes[0].methods).toHaveLength(3)
expect(ast.classes[0].methods[0].decorators).toContain("@Get()")
expect(ast.classes[0].methods[1].decorators).toContain("@Get(':id')")
expect(ast.classes[0].methods[2].decorators).toContain("@Post()")
})
})
})

View File

@@ -1773,4 +1773,244 @@ describe("prompts", () => {
expect(context).toContain("- enum VeryLongEnumName")
})
})
describe("decorators (0.24.4)", () => {
it("should format function with decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["controller.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"controller.ts",
{
imports: [],
exports: [],
functions: [
{
name: "getUser",
lineStart: 1,
lineEnd: 5,
params: [
{
name: "id",
type: "string",
optional: false,
hasDefault: false,
},
],
isAsync: true,
isExported: true,
returnType: "Promise<User>",
decorators: ["@Get(':id')"],
},
],
classes: [],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain("- @Get(':id') async getUser(id: string): Promise<User>")
})
it("should format class with decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["controller.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"controller.ts",
{
imports: [],
exports: [],
functions: [],
classes: [
{
name: "UserController",
lineStart: 1,
lineEnd: 20,
methods: [],
properties: [],
implements: [],
isExported: true,
isAbstract: false,
decorators: ["@Controller('users')"],
},
],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain("- @Controller('users') class UserController")
})
it("should format class with multiple decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["service.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"service.ts",
{
imports: [],
exports: [],
functions: [],
classes: [
{
name: "UserService",
lineStart: 1,
lineEnd: 30,
methods: [],
properties: [],
implements: ["IUserService"],
isExported: true,
isAbstract: false,
decorators: ["@Injectable()", "@Singleton()"],
},
],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain(
"- @Injectable() @Singleton() class UserService implements IUserService",
)
})
it("should format function with multiple decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["controller.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"controller.ts",
{
imports: [],
exports: [],
functions: [
{
name: "createUser",
lineStart: 1,
lineEnd: 10,
params: [],
isAsync: true,
isExported: true,
decorators: ["@Post()", "@Auth()", "@ValidateBody()"],
},
],
classes: [],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain("- @Post() @Auth() @ValidateBody() async createUser()")
})
it("should handle function without decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["utils.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"utils.ts",
{
imports: [],
exports: [],
functions: [
{
name: "helper",
lineStart: 1,
lineEnd: 5,
params: [],
isAsync: false,
isExported: true,
},
],
classes: [],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain("- helper()")
expect(context).not.toContain("@")
})
it("should handle class without decorators", () => {
const structure: ProjectStructure = {
name: "test",
rootPath: "/test",
files: ["simple.ts"],
directories: [],
}
const asts = new Map<string, FileAST>([
[
"simple.ts",
{
imports: [],
exports: [],
functions: [],
classes: [
{
name: "SimpleClass",
lineStart: 1,
lineEnd: 10,
methods: [],
properties: [],
implements: [],
isExported: true,
isAbstract: false,
},
],
interfaces: [],
typeAliases: [],
parseError: false,
},
],
])
const context = buildInitialContext(structure, asts)
expect(context).toContain("- class SimpleClass")
expect(context).not.toContain("@")
})
})
})