Infer missing typescript types

infer.ts

import * as fs from "fs";
import * as ts from "typescript";

export function infer(services: ts.LanguageService, fileNames: string[]) {

    function isNumber(t: ts.Type): boolean {
        return (t.flags & (ts.TypeFlags.Number | ts.TypeFlags.NumberLiteral)) !== 0;
    }
    function isAny(t: ts.Type): boolean {
        return (t.flags & ts.TypeFlags.Any) !== 0;
    }
    function isStructuredType(t: ts.Type): boolean {
        return (t.flags & ts.TypeFlags.StructuredType) !== 0;
    }
    // function getCallSignatures(t: ts.Type): ts.Signature[] {
    //     // if (!isStructuredType(t)) {
    //     //     return [];
    //     // }
    //     console.log(`KEYS: ${Object.keys(t)}`);
    //     return t.getCallSignatures();
    //     // if ('declaredCallSignatures' in t) {
    //     //     const it = <ts.InterfaceTypeWithDeclaredMembers>t;
    //     //     return it.declaredCallSignatures;
    //     // }
    //     return [];
    //     // return (t.flags & ts.TypeFlags.StructuredType) !== 0;
    // }
    
    function typeAndSymbol(n: ts.Node): [ts.Type, ts.Symbol] {
        let symbolNode: ts.Node;
        switch (n.kind) {
            case ts.SyntaxKind.Parameter:
                symbolNode = (<ts.ParameterDeclaration>n).name;
                break;
            case ts.SyntaxKind.VariableDeclaration:
                symbolNode = (<ts.VariableDeclaration>n).name;
                break;
            default:
                symbolNode = n;
        }
        return [checker.getTypeAtLocation(n), checker.getSymbolAtLocation(symbolNode)];
    }

    function typeToString(t: ts.Type | string) {
        return typeof t === 'string' ? t : checker.typeToString(t);
    }

    function fmtType(node: ts.Node) {
        const type = checker.getTypeAtLocation(node);
        return typeToString(type);
    }

    const allConstraints = new Map<ts.Symbol, (ts.Type | string)[]>();

    function addConstraint(sym: ts.Symbol, type: ts.Type | string) {
        console.log(`CONSTRAINT: ${checker.symbolToString(sym)} is a ${typeof type === 'string' ? type : typeToString(type)}!`);

        let constraints = allConstraints.get(sym);
        if (constraints) {
            constraints.push(type);
        } else {
            allConstraints.set(sym, constraints = [type]);
        }
    }

    /** visit nodes finding exported classes */    
    function visit(fileName: string, node: ts.Node) {
        // Only consider exported nodes
        console.log(`[${fileName}] Node: ${node.kind}`);

        const ctxType = checker.getContextualType(<ts.Expression>node);
        if (ctxType && !isAny(ctxType)) {
            const [nodeType, nodeSym] = typeAndSymbol(node);
            if (nodeSym && isAny(nodeType)) {
                // const paramType = checker.getTypeOfSymbolAtLocation(param, node);//param.declarations[0]);
                // addConstraint(nodeSym, paramType);
                addConstraint(nodeSym, ctxType);
            }
        }

        // if (node.kind === ts.SyntaxKind.Parameter) {
        //     const param = <ts.ParameterDeclaration>node;
        //     if (param.name.kind === ts.SyntaxKind.Identifier) {
        //         const name = param.name.text;
        //         const [type, symbol] = typeAndSymbol(param);
        //         console.log(`[${fileName}] Param(${name}): ${typeToString(type)}`);
        //         console.log(`[${fileName}] symbol = ${checker.symbolToString(symbol)}`);
        //         // const aliasedSymbol = checker.getAliasedSymbol(symbol);
        //         // console.log(`[${fileName}] aliased symbol = ${checker.symbolToString(aliasedSymbol)}`);
        //         //     // const symbol = checker.getSymbolsOfParameterPropertyDeclaration(param, name);
        //     }
        // }
        if (node.kind === ts.SyntaxKind.CallExpression) {
            const call = <ts.CallExpression>node;
            if (call.expression.kind === ts.SyntaxKind.PropertyAccessExpression) {
                const access = <ts.PropertyAccessExpression>call.expression;
                const [targetType, targetSym] = typeAndSymbol(access.expression);
                
                if (targetSym && isAny(targetType)) {//} && call.name && call.name.kind == ts.SyntaxKind.Identifier) {
                    const returnType = ctxType;
                    const argTypes = call.arguments.map(a => checker.getTypeAtLocation(a));
                    addConstraint(targetSym, `{\n\t${access.name.text}(${argTypes.map(typeToString).join(', ')}): ${returnType ? typeToString(returnType) : 'any'}\n}`);
                }
            }
        //     // const signatures = targetType.getCallSignatures();
        //     // console.log(`Found ${signatures.length} function signatures for ${typeToString(targetType)}!`);
        //     // if (signatures.length > 0) {
        //     //     const [signature] = signatures;

        //         for (let i = 0; i < call.arguments.length; i++) {
        //             // const param = signature.parameters[i];
        //             const arg = call.arguments[i];
        //             const ctxType = checker.getContextualType(arg);
        //             const [argType, argSym] = typeAndSymbol(arg);
        //             if (argSym && isAny(argType)) {
        //                 // const paramType = checker.getTypeOfSymbolAtLocation(param, arg);//param.declarations[0]);
        //                 // addConstraint(argSym, paramType);
        //                 addConstraint(argSym, ctxType);
        //             }
        //         }
        //         // <ts.Method
        //     // }
            // console.log(`[${fileName}] Call: 
            //     text = "${node.getFullText()}"
            //     target = "${call.expression.getFullText()}"
            //         type = ${typeToString(targetType)}
            //         symbol = ${targetSym ? checker.symbolToString(targetSym) : ''}
            // `);
        }
        if (node.kind === ts.SyntaxKind.BinaryExpression) {
            const binExpr = <ts.BinaryExpression>node;
            const [[leftType, leftSym], [rightType, rightSym]] = [binExpr.left, binExpr.right].map(typeAndSymbol);
            
            console.log(`[${fileName}] Binary expression: ${node.kind}
                text = "${node.getFullText()}"
                left.type = ${typeToString(leftType)} (symbol = ${leftSym ? checker.symbolToString(leftSym) : ''})
                right.type = ${typeToString(rightType)} (symbol = ${rightSym ? checker.symbolToString(rightSym) : ''})
            `);

            if (leftSym || rightType) {
                switch (binExpr.operatorToken.kind) {
                    case ts.SyntaxKind.PlusToken:
                        if (leftSym && isAny(leftType)) {
                            console.log(`left is any`);
                            console.log(`right.flags = ${rightType.flags}`);
                            if (isNumber(rightType)) {
                                addConstraint(leftSym, rightType);
                            }
                        }
                        break;
                    default:
                        console.log(`OP: ${binExpr.operatorToken.kind}`);
                }
            }
        }
        
        ts.forEachChild(node, n => visit(fileName, n));

        // if (node.kind === ts.SyntaxKind.ClassDeclaration) {
        //     // This is a top level class, get its symbol
        //     let symbol = checker.getSymbolAtLocation((<ts.ClassDeclaration>node).name);
        //     output.push(serializeClass(symbol));
        //     // No need to walk any further, class expressions/inner declarations
        //     // cannot be exported
        // }
        // else if (node.kind === ts.SyntaxKind.ModuleDeclaration) {
        //     // This is a namespace, visit its children
        //     ts.forEachChild(node, visit);
        // }
    }

    const program = services.getProgram();
    let checker = program.getTypeChecker();
    // Visit every sourceFile in the program    
    for (const sourceFile of program.getSourceFiles()) {
        // Walk the tree to search for classes
        if (fileNames.indexOf(sourceFile.fileName) >= 0) {
            ts.forEachChild(sourceFile, n => visit(sourceFile.fileName, n));
        }
        // ts.forEachChild(sourceFile, visit);
    }

    for (const [sym, types] of allConstraints) {
        console.log(`CONSTRAINTS for ${checker.symbolToString(sym)}: ${types.map(typeToString)}`);
        let [decl] = sym.getDeclarations();
        console.log(`${decl.getFullText()}: ${types.map(typeToString).join('|')}`)
        // decl.getEnd()
    }
}

input.ts

function f(x, y) {
  console.log(x, y);
  let z = x.call(1, 2);
  g(z);
  return x + 2 + g(y);
}
function g(x: number) {
  return 1;
}

main.ts

/// <reference path="../node_modules/@types/node/index.d.ts" />

import * as fs from "fs";
import * as ts from "typescript";
import {infer} from './infer';

function watch(rootFileNames: string[], options: ts.CompilerOptions) {
    const files: ts.MapLike<{ version: number }> = {};

    // initialize the list of files
    rootFileNames.forEach(fileName => {
        files[fileName] = { version: 0 };
    });

    // Create the language service host to allow the LS to communicate with the host
    const servicesHost: ts.LanguageServiceHost = {
        getScriptFileNames: () => rootFileNames,
        getScriptVersion: (fileName) => files[fileName] && files[fileName].version.toString(),
        getScriptSnapshot: (fileName) => {
            if (!fs.existsSync(fileName)) {
                return undefined;
            }

            return ts.ScriptSnapshot.fromString(fs.readFileSync(fileName).toString());
        },
        getCurrentDirectory: () => process.cwd(),
        getCompilationSettings: () => options,
        getDefaultLibFileName: (options) => ts.getDefaultLibFilePath(options),
    };

    // Create the language service files
    const services = ts.createLanguageService(servicesHost, ts.createDocumentRegistry())

    // Now let's watch the files
    rootFileNames.forEach(fileName => {
        // First time around, emit all files
        emitFile(fileName);

        // Add a watch on the file to handle next change
        fs.watchFile(fileName,
            { persistent: true, interval: 250 },
            (curr, prev) => {
                // Check timestamp
                if (+curr.mtime <= +prev.mtime) {
                    return;
                }

                // Update the version to signal a change in the file
                files[fileName].version++;

                // write the changes to disk
                emitFile(fileName);
            });
    });

    /** visit nodes finding exported classes */    
    function visit(node: ts.Node) {
        // Only consider exported nodes
        console.log(`Node: ${node.kind}`);

        // if (node.kind === ts.SyntaxKind.ClassDeclaration) {
        //     // This is a top level class, get its symbol
        //     let symbol = checker.getSymbolAtLocation((<ts.ClassDeclaration>node).name);
        //     output.push(serializeClass(symbol));
        //     // No need to walk any further, class expressions/inner declarations
        //     // cannot be exported
        // }
        // else if (node.kind === ts.SyntaxKind.ModuleDeclaration) {
        //     // This is a namespace, visit its children
        //     ts.forEachChild(node, visit);
        // }
    }

    function emitFile(fileName: string) {
        logErrors(fileName);
        infer(services, fileNames);

        // let output: DocEntry[] = [];
        // let output = services.getEmitOutput(fileName);

        // if (!output.emitSkipped) {
        //     console.log(`Emitting ${fileName}`);
        // }
        // else {
        //     console.log(`Emitting ${fileName} failed`);
        //     logErrors(fileName);
        // }

        // printProgram();
        // output.outputFiles.forEach(o => {
        //     fs.writeFileSync(o.name, o.text, "utf8");
        // });
    }

    function logErrors(fileName: string) {
        let allDiagnostics = services.getCompilerOptionsDiagnostics()
            .concat(services.getSyntacticDiagnostics(fileName))
            .concat(services.getSemanticDiagnostics(fileName));

        allDiagnostics.forEach(diagnostic => {
            let message = ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n");
            if (diagnostic.file) {
                let { line, character } = diagnostic.file.getLineAndCharacterOfPosition(diagnostic.start);
                console.log(`  Error ${diagnostic.file.fileName} (${line + 1},${character + 1}): ${message}`);
            }
            else {
                console.log(`  Error: ${message}`);
            }
        });
    }
}

// // Initialize files constituting the program as all .ts files in the current directory
// const currentDirectoryFiles = fs.readdirSync(process.cwd()).
//     filter(fileName=> fileName.length >= 3 && fileName.substr(fileName.length - 3, 3) === ".ts");

// // Start the watcher
// watch(currentDirectoryFiles, { module: ts.ModuleKind.CommonJS });

const fileNames = process.argv.slice(2);

// Start the watcher
watch(fileNames, { module: ts.ModuleKind.CommonJS });

package.json

{
  "devDependencies": {
    "concurrently": "^3.1.0",
    "typescript": "^2.1.5",
    "watch": "^1.0.1"
  },
  "name": "infer-types.ts",
  "version": "1.0.0",
  "main": "index.js",
  "dependencies": {
    "@types/acorn": "^4.0.0",
    "@types/estree": "0.0.34",
    "@types/node": "^7.0.4",
    "acorn": "^3.3.0"
  },
  "scripts": {
    "test:w": "tsc && concurrently \"tsc -w\" \"watch --wait=1 \\\"node build/src/main2.js test/data/input.ts\\\" build test/data\""
  },
  "author": "",
  "license": "ISC"
}

tsconfig.json

{
    "compilerOptions": {
        "target": "es6",
        "module": "commonjs",
        "moduleResolution": "node",
        "outDir": "build",
        "forceConsistentCasingInFileNames": true,
        "noImplicitAny": false,
        "noImplicitReturns": true,
        "noFallthroughCasesInSwitch": true,
        "pretty": true,
        "removeComments": true,
        "sourceMap": true,
        "strictNullChecks": true
    },
    "atom": {
        "rewriteTsconfig": true
    }
}