@@ -86,8 +86,28 @@ module.exports = {
8686 }
8787 }
8888
89- function getReturnStatement ( body ) {
90- return body . filter ( ( item ) => item . type === 'ReturnStatement' ) [ 0 ] ;
89+ function getReturnStatements ( node ) {
90+ const returnStatements = arguments [ 1 ] || [ ] ;
91+ if ( node . type === 'IfStatement' ) {
92+ if ( node . consequent ) {
93+ getReturnStatements ( node . consequent , returnStatements ) ;
94+ }
95+ if ( node . alternate ) {
96+ getReturnStatements ( node . alternate , returnStatements ) ;
97+ }
98+ } else {
99+ node . body . forEach ( ( item ) => {
100+ if ( item . type === 'IfStatement' ) {
101+ getReturnStatements ( item , returnStatements ) ;
102+ }
103+
104+ if ( item . type === 'ReturnStatement' ) {
105+ returnStatements . push ( item ) ;
106+ }
107+ } ) ;
108+ }
109+
110+ return returnStatements ;
91111 }
92112
93113 function isKeyAfterSpread ( attributes ) {
@@ -189,10 +209,11 @@ module.exports = {
189209
190210 if ( isFn || isArrFn ) {
191211 if ( fn . body . type === 'BlockStatement' ) {
192- const returnStatement = getReturnStatement ( fn . body . body ) ;
193- if ( returnStatement && returnStatement . argument ) {
194- checkIteratorElement ( returnStatement . argument ) ;
195- }
212+ getReturnStatements ( fn . body )
213+ . filter ( ( returnStatement ) => returnStatement && returnStatement . argument )
214+ . forEach ( ( returnStatement ) => {
215+ checkIteratorElement ( returnStatement . argument ) ;
216+ } ) ;
196217 }
197218 }
198219 } ,
0 commit comments