find return statements in function body with typescript compiler api

Followup of Find node of specific SyntaxKind in ASL tree with typescript compiler api but focused on return statements for future googlers.

I would like to find return statements in a method body with the typescript compiler api.

My function body is unknown, so I hope to not have to access the different nodes statically. I found there is a function visitFunctionBody but that requires a context. The example I've seen seem to get a context from a transformer factory but typescript does not export such a member anymore and I'm not even sure this does what I need to do.

So how can one find return statements in a body ?

I've tried something along the lines of:

function findReturnStatements() {
    const children: Node[] = [];
    node.forEachChild((childNode) => {
      children.push(childNode);
    });
    return node.filter(isReturnStatement);
}

But this can fail depending on the shape of the body

someMethod() {
  return 3;
}

// vs 

someMethod() {
  try {
    return 3;
  } catch (e) {
    return 4;
  }
}

You probably want something like this:

// untested and very likely has a bug, but it should be a good starting point
function* getReturnStmts(node: ts.Node): Iterable<ts.ReturnStatement | ts.Expression> {
  if (ts.isReturnStatement(node)) {
    yield node;
  } else if (ts.isBlock(node) || ts.isCaseClause(node) || ts.isDefaultClause(node)) {
    for (const stmt of node.statements) {
      yield* getReturnStmts(stmt);
    }
  } else if (ts.isIfStatement(node)) {
    yield* getReturnStmts(node.thenStatement);
    if (node.elseStatement) {
      yield* getReturnStmts(node.elseStatement);
    }
  } else if (ts.isIterationStatement(node, true)) {
    yield* getReturnStmts(node.statement);
  } else if (ts.isSwitchStatement(node)) {
    for (const clause of node.caseBlock.clauses) {
      yield* getReturnStmts(clause);
    }
  } else if (ts.isTryStatement(node)) {
    yield* getReturnStmts(node.tryBlock);
    if (node.catchClause) {
      yield* getReturnStmts(node.catchClause.block);
    }
    if (node.finallyBlock) {
      yield* getReturnStmts(node.finallyBlock);
    }
  } else if (
    ts.isMethodDeclaration(node) || ts.isFunctionDeclaration(node) ||
    ts.isArrowFunction(node) || ts.isFunctionExpression(node) ||
    ts.isConstructorDeclaration(node)
  ) {
    if (node.body != null) {
      if (ts.isBlock(node.body)) {
        yield* getReturnStmts(node.body);
      } else {
        yield node.body;
      }
    }
  }
}

Note that arrow functions can have a consise body, which is just an expression (ex. () => 1), so that's why it returns ts.Expression in addition to ts.ReturnStatement.