diff --git a/packages/babel-traverse/src/path/replacement.ts b/packages/babel-traverse/src/path/replacement.ts index e114740501e8..e0e4f5bcc8e4 100644 --- a/packages/babel-traverse/src/path/replacement.ts +++ b/packages/babel-traverse/src/path/replacement.ts @@ -228,6 +228,7 @@ export function replaceExpressionWithStatements( const functionParent = this.getFunctionParent(); const isParentAsync = functionParent?.is("async"); + const isParentGenerator = functionParent?.is("generator"); const container = t.arrowFunctionExpression([], t.blockStatement(nodes)); @@ -277,17 +278,28 @@ export function replaceExpressionWithStatements( callee.arrowFunctionToExpression(); // (() => await xxx)() -> await (async () => await xxx)(); - if ( + const needToAwaitFunction = isParentAsync && traverse.hasType( (this.get("callee.body") as NodePath).node, "AwaitExpression", t.FUNCTION_TYPES, - ) - ) { + ); + const needToYieldFunction = + isParentGenerator && + traverse.hasType( + (this.get("callee.body") as NodePath).node, + "YieldExpression", + t.FUNCTION_TYPES, + ); + if (needToAwaitFunction) { callee.set("async", true); this.replaceWith(t.awaitExpression((this as ThisType).node)); } + if (needToYieldFunction) { + callee.set("generator", true); + this.replaceWith(t.yieldExpression((this as ThisType).node, true)); + } return callee.get("body.body"); }