diff --git a/src/exports-symbol-tree.ts b/src/exports-symbol-tree.ts index 34dd656..19199a4 100644 --- a/src/exports-symbol-tree.ts +++ b/src/exports-symbol-tree.ts @@ -58,12 +58,13 @@ export class ExportsSymbolTree { } private computeTreeForChildren(targetSymbolsSet: Set, node: ts.Node, visitedSymbols: Set): void { + const typeChecker = this.program.getTypeChecker(); + // it's similar to handling ts.Block node - both Block and variable's initializer are part of _implementation_ // and we don't care about that implementation at all - we just only need to worry it's definition // for functions it is arguments and return type // for variables - the type of a variable if (ts.isVariableDeclaration(node)) { - const typeChecker = this.program.getTypeChecker(); const variableType = typeChecker.getTypeAtLocation(node); const variableTypeSymbol = variableType.getSymbol(); if (variableTypeSymbol !== undefined) { @@ -73,6 +74,15 @@ export class ExportsSymbolTree { return; } + if ((ts.isMethodDeclaration(node) || ts.isFunctionDeclaration(node)) && node.type === undefined) { + // this means that a function/method doesn't have a declared returned type so here we need to get calculated one + for (const signature of typeChecker.getTypeAtLocation(node).getCallSignatures()) { + for (const childSymbol of splitTransientSymbol(signature.getReturnType().symbol, typeChecker)) { + targetSymbolsSet.add(childSymbol); + } + } + } + ts.forEachChild(node, (childNode: ts.Node) => this.computeTreeForNode(targetSymbolsSet, childNode, visitedSymbols)); } diff --git a/src/transformer.ts b/src/transformer.ts index ba7132f..1756059 100644 --- a/src/transformer.ts +++ b/src/transformer.ts @@ -437,10 +437,19 @@ function createTransformerFactory(program: ts.Program, options?: Partial