@@ -83,7 +83,7 @@ export function getSetupFunction(
83
83
}
84
84
85
85
/**
86
- * Extracts variable and function declarations from a function body
86
+ * Extracts variable and function declarations from a function body (recursive)
87
87
*/
88
88
export function extractDeclarations ( body : TSESTree . BlockStatement ) : {
89
89
variables : string [ ]
@@ -92,16 +92,65 @@ export function extractDeclarations(body: TSESTree.BlockStatement): {
92
92
const variables : string [ ] = [ ]
93
93
const functions : string [ ] = [ ]
94
94
95
- for ( const statement of body . body ) {
96
- if ( statement . type === 'VariableDeclaration' ) {
97
- for ( const declarator of statement . declarations ) {
98
- extractIdentifiersFromPattern ( declarator . id , variables )
99
- }
100
- } else if ( statement . type === 'FunctionDeclaration' && statement . id ) {
101
- functions . push ( statement . id . name )
95
+ function traverse ( node : TSESTree . Node ) : void {
96
+ switch ( node . type ) {
97
+ case 'VariableDeclaration' :
98
+ for ( const declarator of node . declarations ) {
99
+ extractIdentifiersFromPattern ( declarator . id , variables )
100
+ }
101
+ break
102
+ case 'FunctionDeclaration' :
103
+ if ( node . id ) {
104
+ functions . push ( node . id . name )
105
+ }
106
+ break
107
+ case 'BlockStatement' :
108
+ for ( const statement of node . body ) {
109
+ traverse ( statement )
110
+ }
111
+ break
112
+ case 'IfStatement' :
113
+ traverse ( node . consequent )
114
+ if ( node . alternate ) {
115
+ traverse ( node . alternate )
116
+ }
117
+ break
118
+ case 'ForStatement' :
119
+ case 'ForInStatement' :
120
+ case 'ForOfStatement' :
121
+ if ( node . body ) {
122
+ traverse ( node . body )
123
+ }
124
+ break
125
+ case 'WhileStatement' :
126
+ case 'DoWhileStatement' :
127
+ traverse ( node . body )
128
+ break
129
+ case 'SwitchStatement' :
130
+ for ( const switchCase of node . cases ) {
131
+ for ( const statement of switchCase . consequent ) {
132
+ traverse ( statement )
133
+ }
134
+ }
135
+ break
136
+ case 'TryStatement' :
137
+ traverse ( node . block )
138
+ if ( node . handler ) {
139
+ traverse ( node . handler . body )
140
+ }
141
+ if ( node . finalizer ) {
142
+ traverse ( node . finalizer )
143
+ }
144
+ break
145
+ case 'WithStatement' :
146
+ traverse ( node . body )
147
+ break
148
+ // For other statement types, we don't need to traverse deeper
149
+ // as they don't contain variable/function declarations
102
150
}
103
151
}
104
152
153
+ traverse ( body )
105
154
return { variables, functions }
106
155
}
107
156
@@ -219,15 +268,102 @@ export function hasSpreadInReturn(
219
268
}
220
269
221
270
/**
222
- * Finds the return statement in a function body
271
+ * Finds all return statements in a function body (recursive)
272
+ */
273
+ export function findAllReturnStatements (
274
+ body : TSESTree . BlockStatement
275
+ ) : TSESTree . ReturnStatement [ ] {
276
+ const returnStatements : TSESTree . ReturnStatement [ ] = [ ]
277
+
278
+ function traverse ( node : TSESTree . Node ) : void {
279
+ switch ( node . type ) {
280
+ case 'ReturnStatement' :
281
+ returnStatements . push ( node )
282
+ break
283
+ case 'BlockStatement' :
284
+ for ( const statement of node . body ) {
285
+ traverse ( statement )
286
+ }
287
+ break
288
+ case 'IfStatement' :
289
+ traverse ( node . consequent )
290
+ if ( node . alternate ) {
291
+ traverse ( node . alternate )
292
+ }
293
+ break
294
+ case 'ForStatement' :
295
+ case 'ForInStatement' :
296
+ case 'ForOfStatement' :
297
+ if ( node . body ) {
298
+ traverse ( node . body )
299
+ }
300
+ break
301
+ case 'WhileStatement' :
302
+ case 'DoWhileStatement' :
303
+ traverse ( node . body )
304
+ break
305
+ case 'SwitchStatement' :
306
+ for ( const switchCase of node . cases ) {
307
+ for ( const statement of switchCase . consequent ) {
308
+ traverse ( statement )
309
+ }
310
+ }
311
+ break
312
+ case 'TryStatement' :
313
+ traverse ( node . block )
314
+ if ( node . handler ) {
315
+ traverse ( node . handler . body )
316
+ }
317
+ if ( node . finalizer ) {
318
+ traverse ( node . finalizer )
319
+ }
320
+ break
321
+ case 'WithStatement' :
322
+ traverse ( node . body )
323
+ break
324
+ // For function declarations/expressions, we don't traverse into them
325
+ // as they have their own scope
326
+ case 'FunctionDeclaration' :
327
+ case 'FunctionExpression' :
328
+ case 'ArrowFunctionExpression' :
329
+ break
330
+ // For other statement types that can contain nested statements
331
+ case 'ExpressionStatement' :
332
+ case 'VariableDeclaration' :
333
+ case 'ThrowStatement' :
334
+ case 'BreakStatement' :
335
+ case 'ContinueStatement' :
336
+ case 'EmptyStatement' :
337
+ case 'DebuggerStatement' :
338
+ // These don't contain nested statements
339
+ break
340
+ }
341
+ }
342
+
343
+ traverse ( body )
344
+ return returnStatements
345
+ }
346
+
347
+ /**
348
+ * Finds the main return statement in a function body (typically the last object return)
223
349
*/
224
350
export function findReturnStatement (
225
351
body : TSESTree . BlockStatement
226
352
) : TSESTree . ReturnStatement | null {
227
- for ( const statement of body . body ) {
228
- if ( statement . type === 'ReturnStatement' ) {
229
- return statement
353
+ const allReturns = findAllReturnStatements ( body )
354
+
355
+ if ( allReturns . length === 0 ) {
356
+ return null
357
+ }
358
+
359
+ // Find the last return statement that returns an object expression
360
+ for ( let i = allReturns . length - 1 ; i >= 0 ; i -- ) {
361
+ const returnStmt = allReturns [ i ]
362
+ if ( returnStmt . argument ?. type === 'ObjectExpression' ) {
363
+ return returnStmt
230
364
}
231
365
}
232
- return null
366
+
367
+ // If no object return found, return the last return statement
368
+ return allReturns [ allReturns . length - 1 ]
233
369
}
0 commit comments