Find more types using Type Inference

src/apply-types.spec.ts

@@ -0,0 +1,32 @@
+import * as ts from 'typescript';
+import { applyTypes } from './apply-types';
+
+describe('applyTypes', () => {
+    it('should throw an error if given non-existing tsconfig.json file', () => {
+        expect(() => applyTypes([], { tsConfig: 'not-found-file.json' })).toThrowError(
+            `Error while reading not-found-file.json: The specified path does not exist: 'not-found-file.json'.`,
+        );
+    });
+
+    it('should throw an error if given bad tsconfig.json file', () => {
+        const tsConfigHost = {
+            ...ts.sys,
+            readFile: jest.fn(() => '<invalid json>'),
+        };
+        expect(() => applyTypes([], { tsConfig: 'tsconfig.bad.json', tsConfigHost })).toThrowError(
+            `Error while reading tsconfig.bad.json: '{' expected.`,
+        );
+        expect(tsConfigHost.readFile).toHaveBeenCalledWith('tsconfig.bad.json');
+    });
+
+    it('should throw an error in case of validation errors in given tsconfig.json file', () => {
+        const tsConfigHost = {
+            ...ts.sys,
+            readFile: jest.fn(() => '{ "include": 123 }'),
+        };
+        expect(() => applyTypes([], { tsConfig: 'tsconfig.invalid.json', tsConfigHost })).toThrowError(
+            `Error while parsing tsconfig.invalid.json: Compiler option 'include' requires a value of type Array.`,
+        );
+        expect(tsConfigHost.readFile).toHaveBeenCalledWith('tsconfig.invalid.json');
+    });
+});

src/apply-types.ts

@@ -1,10 +1,14 @@
 import * as fs from 'fs';
 import * as path from 'path';
+import * as ts from 'typescript';
 
 import { IExtraOptions } from './instrument';
 import { applyReplacements, Replacement } from './replacement';
+import { ISourceLocation } from './type-collector-snippet';
 
-export type ICollectedTypeInfo = Array<[string, number, string[], IExtraOptions]>;
+export type ICollectedTypeInfo = Array<
+    [string, number, Array<[string | undefined, ISourceLocation | undefined]>, IExtraOptions]
+>;
 
 export interface IApplyTypesOptions {
     /**
@@ -18,20 +22,60 @@ export interface IApplyTypesOptions {
      * If given, all the file paths in the collected type info will be resolved relative to this directory.
      */
     rootDir?: string;
+
+    /**
+     * Path to your project's tsconfig file
+     */
+    tsConfig?: string;
+
+    // You probably never need to touch these two - they are used by the integration tests to setup
+    // a virtual file system for TS:
+    tsConfigHost?: ts.ParseConfigHost;
+    tsCompilerHost?: ts.CompilerHost;
+}
+
+function findType(program?: ts.Program, typeName?: string, sourcePos?: ISourceLocation) {
+    if (program && sourcePos) {
+        const [sourceName, sourceOffset] = sourcePos;
+        const typeChecker = program.getTypeChecker();
+        let foundType: string | null = null;
+        function visit(node: ts.Node) {
+            if (node.getStart() === sourceOffset) {
+                const type = typeChecker.getTypeAtLocation(node);
+                foundType = typeChecker.typeToString(type);
+            }
+            ts.forEachChild(node, visit);
+        }
+        const sourceFile = program.getSourceFile(sourceName);
+        visit(sourceFile);
+        if (foundType) {
+            return foundType;
+        }
+    }
+    return typeName;
 }
 
-export function applyTypesToFile(source: string, typeInfo: ICollectedTypeInfo, options: IApplyTypesOptions) {
+export function applyTypesToFile(
+    source: string,
+    typeInfo: ICollectedTypeInfo,
+    options: IApplyTypesOptions,
+    program?: ts.Program,
+) {
     const replacements = [];
     const prefix = options.prefix || '';
     for (const [, pos, types, opts] of typeInfo) {
         const isOptional = source[pos - 1] === '?';
-        let sortedTypes = types.sort();
+        let sortedTypes = types
+            .map(([name, sourcePos]) => findType(program, name, sourcePos))
+            .filter((t) => t)
+            .sort();
         if (isOptional) {
             sortedTypes = sortedTypes.filter((t) => t !== 'undefined');
-            if (sortedTypes.length === 0) {
-                continue;
-            }
         }
+        if (sortedTypes.length === 0) {
+            continue;
+        }
+
         let suffix = '';
         if (opts && opts.parens) {
             replacements.push(Replacement.insert(opts.parens[0], '('));
@@ -44,6 +88,22 @@ export function applyTypesToFile(source: string, typeInfo: ICollectedTypeInfo, o
 
 export function applyTypes(typeInfo: ICollectedTypeInfo, options: IApplyTypesOptions = {}) {
     const files: { [key: string]: typeof typeInfo } = {};
+    let program: ts.Program | undefined;
+    if (options.tsConfig) {
+        const configHost = options.tsConfigHost || ts.sys;
+        const { config, error } = ts.readConfigFile(options.tsConfig, configHost.readFile);
+        if (error) {
+            throw new Error(`Error while reading ${options.tsConfig}: ${error.messageText}`);
+        }
+
+        const parsed = ts.parseJsonConfigFileContent(config, configHost, options.rootDir || '');
+        if (parsed.errors.length) {
+            const errors = parsed.errors.map((e) => e.messageText).join(', ');
+            throw new Error(`Error while parsing ${options.tsConfig}: ${errors}`);
+        }
+
+        program = ts.createProgram(parsed.fileNames, parsed.options, options.tsCompilerHost);
+    }
     for (const entry of typeInfo) {
         const file = entry[0];
         if (!files[file]) {
@@ -54,6 +114,6 @@ export function applyTypes(typeInfo: ICollectedTypeInfo, options: IApplyTypesOpt
     for (const file of Object.keys(files)) {
         const filePath = options.rootDir ? path.join(options.rootDir, file) : file;
         const source = fs.readFileSync(filePath, 'utf-8');
-        fs.writeFileSync(filePath, applyTypesToFile(source, files[file], options));
+        fs.writeFileSync(filePath, applyTypesToFile(source, files[file], options, program));
     }
 }

src/instrument.ts

@@ -2,6 +2,10 @@ import * as ts from 'typescript';
 
 import { applyReplacements, Replacement } from './replacement';
 
+export interface IInstrumentOptions {
+    instrumentCallExpressions: boolean;
+}
+
 export interface IExtraOptions {
     arrow?: boolean;
     parens?: [number, number];
@@ -21,7 +25,7 @@ function hasParensAroundArguments(node: ts.FunctionLike) {
     }
 }
 
-function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
+function visit(node: ts.Node, replacements: Replacement[], fileName: string, options: IInstrumentOptions) {
     const isArrow = ts.isArrowFunction(node);
     if (ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node) || ts.isArrowFunction(node)) {
         const isShortArrow = ts.isArrowFunction(node) && !ts.isBlock(node.body);
@@ -45,14 +49,27 @@ function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
                 const instrumentExpr = `$_$twiz(${params.join(',')})`;
                 if (isShortArrow) {
                     replacements.push(Replacement.insert(node.body.getStart(), `(${instrumentExpr},`));
-                    replacements.push(Replacement.insert(node.body.getEnd(), `)`));
+                    replacements.push(Replacement.insert(node.body.getEnd(), `)`, 10));
                 } else {
                     replacements.push(Replacement.insert(node.body.getStart() + 1, `${instrumentExpr};`));
                 }
             }
         }
     }
 
+    if (
+        options.instrumentCallExpressions &&
+        ts.isCallExpression(node) &&
+        node.expression.getText() !== 'require.context'
+    ) {
+        for (const arg of node.arguments) {
+            if (!ts.isStringLiteral(arg) && !ts.isNumericLiteral(arg)) {
+                replacements.push(Replacement.insert(arg.getStart(), '$_$twiz.track('));
+                replacements.push(Replacement.insert(arg.getEnd(), `,${JSON.stringify(fileName)},${arg.getStart()})`));
+            }
+        }
+    }
+
     if (
         ts.isPropertyDeclaration(node) &&
         ts.isIdentifier(node.name) &&
@@ -88,16 +105,24 @@ function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
         }
     }
 
-    node.forEachChild((child) => visit(child, replacements, fileName));
+    node.forEachChild((child) => visit(child, replacements, fileName, options));
 }
 
-const declaration =
-    'declare function $_$twiz(name: string, value: any, pos: number, filename: string, opts: any): void;\n';
+const declaration = `
+    declare function $_$twiz(name: string, value: any, pos: number, filename: string, opts: any): void;
+    declare namespace $_$twiz {
+        function track<T>(v: T, p: number, f: string): T;
+    }
+`;
 
-export function instrument(source: string, fileName: string) {
+export function instrument(source: string, fileName: string, options?: IInstrumentOptions) {
+    const instrumentOptions: IInstrumentOptions = {
+        instrumentCallExpressions: false,
+        ...options,
+    };
     const sourceFile = ts.createSourceFile(fileName, source, ts.ScriptTarget.Latest, true);
     const replacements = [] as Replacement[];
-    visit(sourceFile, replacements, fileName);
+    visit(sourceFile, replacements, fileName, instrumentOptions);
     if (replacements.length) {
         replacements.push(Replacement.insert(0, declaration));
     }

src/integration.spec.ts

@@ -2,7 +2,7 @@ import * as fs from 'fs';
 import * as ts from 'typescript';
 import * as vm from 'vm';
 
-import { transpileSource } from './test-utils/transpile';
+import { transpileSource, virtualCompilerHost } from './test-utils/transpile';
 
 const mockFs = {
     readFileSync: jest.fn(fs.readFileSync),
@@ -14,7 +14,7 @@ import { applyTypes, getTypeCollectorSnippet, IApplyTypesOptions, instrument } f
 
 function typeWiz(input: string, typeCheck = false, options?: IApplyTypesOptions) {
     // Step 1: instrument the source
-    const instrumented = instrument(input, 'c:\\test.ts');
+    const instrumented = instrument(input, 'c:\\test.ts', { instrumentCallExpressions: true });
 
     // Step 2: compile + add the type collector
     const compiled = typeCheck ? transpileSource(instrumented, 'test.ts') : ts.transpile(instrumented);
@@ -26,8 +26,35 @@ function typeWiz(input: string, typeCheck = false, options?: IApplyTypesOptions)
     mockFs.readFileSync.mockReturnValue(input);
     mockFs.writeFileSync.mockImplementationOnce(() => 0);
 
+    if (options && options.tsConfig) {
+        options.tsCompilerHost = virtualCompilerHost(input, 'c:/test.ts');
+        options.tsConfigHost = {
+            fileExists: ts.sys.fileExists,
+
+            // readDirectory will be called by applyTypes to get the names of files included in
+            // the typescript project. We return a mock value with our test script path.
+            readDirectory: () => ['c:/test.ts'],
+
+            // readFile will be called to read the compiler options from tsconfig.json, so we mock
+            // it to return a basic configuration that will be used during the integration tests
+            readFile: jest.fn(() =>
+                JSON.stringify({
+                    compilerOptions: {
+                        target: 'es2015',
+                    },
+                    include: ['test.ts'],
+                }),
+            ),
+            useCaseSensitiveFileNames: ts.sys.useCaseSensitiveFileNames,
+        };
+    }
+
     applyTypes(collectedTypes, options);
 
+    if (options && options.tsConfig) {
+        expect(options.tsConfigHost.readFile).toHaveBeenCalledWith(options.tsConfig);
+    }
+
     if (mockFs.writeFileSync.mock.calls.length) {
         expect(mockFs.writeFileSync).toHaveBeenCalledTimes(1);
         expect(mockFs.writeFileSync).toHaveBeenCalledWith('c:\\test.ts', expect.any(String));
@@ -173,6 +200,44 @@ describe('function parameters', () => {
             optional() + optional(10);
         `);
     });
+
+    it('should use TypeScript inference to find argument types', () => {
+        const input = `
+            function f(a) {
+            }
+
+            const arr: string[] = [];
+            f(arr);
+        `;
+
+        expect(typeWiz(input, false, { tsConfig: 'tsconfig.integration.json' })).toBe(`
+            function f(a: string[]) {
+            }
+
+            const arr: string[] = [];
+            f(arr);
+        `);
+    });
+
+    it('should discover generic types using Type Inference', () => {
+        const input = `
+            function f(a) {
+                return a;
+            }
+
+            const promise = Promise.resolve(15);
+            f(promise);
+        `;
+
+        expect(typeWiz(input, false, { tsConfig: 'tsconfig.integration.json' })).toBe(`
+            function f(a: Promise<number>) {
+                return a;
+            }
+
+            const promise = Promise.resolve(15);
+            f(promise);
+        `);
+    });
 });
 
 describe('class fields', () => {

src/replacement.ts

@@ -1,17 +1,22 @@
 export class Replacement {
-    public static insert(pos: number, text: string) {
-        return new Replacement(pos, pos, text);
+    public static insert(pos: number, text: string, priority = 0) {
+        return new Replacement(pos, pos, text, priority);
     }
 
     public static delete(start: number, end: number) {
         return new Replacement(start, end, '');
     }
 
-    constructor(readonly start: number, readonly end: number, readonly text = '') {}
+    constructor(readonly start: number, readonly end: number, readonly text = '', readonly priority = 0) {}
 }
 
 export function applyReplacements(source: string, replacements: Replacement[]) {
-    replacements = replacements.sort((r1, r2) => (r2.end !== r1.end ? r2.end - r1.end : r2.start - r1.start));
+    replacements = replacements.sort(
+        (r1, r2) =>
+            r2.end !== r1.end
+                ? r2.end - r1.end
+                : r1.start !== r2.start ? r2.start - r1.start : r1.priority - r2.priority,
+    );
     for (const replacement of replacements) {
         source = source.slice(0, replacement.start) + replacement.text + source.slice(replacement.end);
     }

src/test-utils/transpile.ts

@@ -1,20 +1,26 @@
 import * as ts from 'typescript';
 
-// similar to ts.transpile(), but also does type checking and throws in case of error
-export function transpileSource(input: string, filename: string) {
-    const compilerOptions = {
-        target: ts.ScriptTarget.ES2015,
-    };
-
+export function virtualCompilerHost(input: string, filename: string, compilerOptions: ts.CompilerOptions = {}) {
     const host = ts.createCompilerHost(compilerOptions);
     const old = host.getSourceFile;
-    host.getSourceFile = (name: string, target: ts.ScriptTarget, ...args) => {
+    host.getSourceFile = (name: string, target: ts.ScriptTarget, ...args: any[]) => {
         if (name === filename) {
             return ts.createSourceFile(filename, input, target, true);
         }
         return old.call(host, name, target, ...args);
     };
 
+    return host;
+}
+
+// similar to ts.transpile(), but also does type checking and throws in case of error
+export function transpileSource(input: string, filename: string) {
+    const compilerOptions = {
+        target: ts.ScriptTarget.ES2015,
+    };
+
+    const host = virtualCompilerHost(input, filename, compilerOptions);
+
     let outputText;
     host.writeFile = (name: string, value: string) => {
         if (name.endsWith('.js')) {