find return statements in function body with typescript compiler api

Followup of https://stackoverflow.com/questions/76505561/find-node-of-specific-syntaxkind-in-asl-tree 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.