feat: arrow function support

mockdeep · Jan 18, 2018 · 5077bc4 · 5077bc4
1 parent ac21da4
commit 5077bc4
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 15 deletions.
diff --git a/src/apply-types.ts b/src/apply-types.ts
@@ -1,13 +1,14 @@
 import * as fs from 'fs';
 import * as path from 'path';
 
+import { IExtraOptions } from './instrument';
 import { applyReplacements, Replacement } from './replacement';
 
-export type ICollectedTypeInfo = Array<[string, number, string[]]>;
+export type ICollectedTypeInfo = Array<[string, number, string[], IExtraOptions]>;
 
 export function applyTypesToFile(source: string, typeInfo: ICollectedTypeInfo) {
     const replacements = [];
-    for (const [, pos, types] of typeInfo) {
+    for (const [, pos, types, opts] of typeInfo) {
         const isOptional = source[pos - 1] === '?';
         let sortedTypes = types.sort();
         if (isOptional) {
@@ -16,12 +17,16 @@ export function applyTypesToFile(source: string, typeInfo: ICollectedTypeInfo) {
                 continue;
             }
         }
+        if (opts && opts.parens) {
+            replacements.push(Replacement.insert(opts.parens[0], '('));
+            replacements.push(Replacement.insert(opts.parens[1], ')'));
+        }
         replacements.push(Replacement.insert(pos, ': ' + sortedTypes.join('|')));
     }
     return applyReplacements(source, replacements);
 }
 
-export function applyTypes(typeInfo: Array<[string, number, string[]]>, rootDir?: string) {
+export function applyTypes(typeInfo: ICollectedTypeInfo, rootDir?: string) {
     const files: {[key: string]: typeof typeInfo} = {};
     for (const entry of typeInfo) {
         const file = entry[0];

diff --git a/src/instrument.spec.ts b/src/instrument.spec.ts
@@ -3,17 +3,18 @@ import { instrument } from './instrument';
 describe('instrument', () => {
     it('should instrument function parameters without types', () => {
         const input = `function (a) { return 5; }`;
-        expect(instrument(input, 'test.ts')).toMatch(`function (a) {$_$twiz("a",a,11,"test.ts"); return 5; }`);
+        expect(instrument(input, 'test.ts')).toMatch(`function (a) {$_$twiz("a",a,11,"test.ts",{}); return 5; }`);
     });
 
     it('should correctly instrument optional function parameters', () => {
         const input = `function (a?) { return 5; }`;
-        expect(instrument(input, 'test.ts')).toMatch(`function (a?) {$_$twiz("a",a,12,"test.ts"); return 5; }`);
+        expect(instrument(input, 'test.ts')).toMatch(`function (a?) {$_$twiz("a",a,12,"test.ts",{}); return 5; }`);
     });
 
     it('should instrument class method parameters', () => {
         const input = `class Foo { bar(a) { return 5; } }`;
-        expect(instrument(input, 'test.ts')).toMatch(`class Foo { bar(a) {$_$twiz("a",a,17,"test.ts"); return 5; } }`);
+        expect(instrument(input, 'test.ts'))
+            .toMatch(`class Foo { bar(a) {$_$twiz("a",a,17,"test.ts",{}); return 5; } }`);
     });
 
     it('should not instrument function parameters that already have a type', () => {

diff --git a/src/instrument.ts b/src/instrument.ts
@@ -2,19 +2,48 @@ import * as ts from 'typescript';
 
 import { applyReplacements, Replacement } from './replacement';
 
+export interface IExtraOptions {
+    arrow?: boolean;
+    parens?: [number, number];
+}
+
+function hasParensAroundArguments(node: ts.FunctionLike) {
+    if (ts.isArrowFunction(node)) {
+        return node.parameters.length !== 1
+            || node.getText().substr(0, node.equalsGreaterThanToken.getStart() - node.getStart()).includes('(');
+    } else {
+        return true;
+    }
+}
+
 function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
-    if ((ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node))) {
+    const isArrow = ts.isArrowFunction(node);
+    if (ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node) || ts.isArrowFunction(node)) {
+        const isShortArrow = ts.isArrowFunction(node) && !ts.isBlock(node.body);
         for (const param of node.parameters) {
             if (!param.type && !param.initializer && node.body) {
                 const typeInsertionPos = param.name.getEnd() + (param.questionToken ? 1 : 0);
+                const opts: IExtraOptions = {};
+                if (isArrow) {
+                    opts.arrow = true;
+                }
+                if ( !hasParensAroundArguments(node)) {
+                    opts.parens = [node.parameters[0].getStart(), node.parameters[0].getEnd()];
+                }
                 const params = [
                     JSON.stringify(param.name.getText()),
                     param.name.getText(),
                     typeInsertionPos,
                     JSON.stringify(fileName),
+                    JSON.stringify(opts),
                 ];
-                const instrumentExpr = `$_$twiz(${params.join(',')});`;
-                replacements.push(Replacement.insert(node.body.getStart() + 1, instrumentExpr));
+                const instrumentExpr = `$_$twiz(${params.join(',')})`;
+                if (isShortArrow) {
+                    replacements.push(Replacement.insert(node.body.getStart(), `(${instrumentExpr},`));
+                    replacements.push(Replacement.insert(node.body.getEnd(), `)`));
+                } else {
+                    replacements.push(Replacement.insert(node.body.getStart() + 1, `${instrumentExpr};`));
+                }
             }
         }
     }
@@ -27,6 +56,7 @@ function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
             'value',
             node.name.getEnd() + (node.questionToken ? 1 : 0),
             JSON.stringify(fileName),
+            JSON.stringify({}),
         ];
         const instrumentExpr = `$_$twiz(${params.join(',')});`;
         const preamble = `
@@ -51,7 +81,8 @@ function visit(node: ts.Node, replacements: Replacement[], fileName: string) {
     node.forEachChild((child) => visit(child, replacements, fileName));
 }
 
-const declaration = `declare function $_$twiz(name: string, value: any, pos: number, filename: string): void;\n`;
+const declaration =
+    `declare function $_$twiz(name: string, value: any, pos: number, filename: string, opts: any): void;\n`;
 
 export function instrument(source: string, fileName: string) {
     const sourceFile = ts.createSourceFile(fileName, source, ts.ScriptTarget.Latest, true);

diff --git a/src/integration.spec.ts b/src/integration.spec.ts
@@ -20,7 +20,7 @@ function typeWiz(input: string, typeCheck = false) {
     const compiled = typeCheck ? transpileSource(instrumented, 'test.ts') : ts.transpile(instrumented);
 
     // Step 3: evaluate the code, and collect the runtime type information
-    const collectedTypes = vm.runInNewContext(getTypeCollectorSnippet() + compiled + '$_$twiz.get();');
+    const collectedTypes = vm.runInNewContext(getTypeCollectorSnippet() + compiled + ';$_$twiz.get();');
 
     // Step 4: put the collected typed into the code
     mockFs.readFileSync.mockReturnValue(input);
@@ -59,6 +59,22 @@ describe('function parameters', () => {
         `);
     });
 
+    it('should find argument types for arrow functions', () => {
+        const input = `((x) => { return x + 5; })(10)`;
+
+        expect(typeWiz(input)).toBe(`((x: number) => { return x + 5; })(10)`);
+    });
+
+    it('should find argument types for arrow function expressions', () => {
+        const input = `((x)=>x+5)(10)`;
+        expect(typeWiz(input)).toBe(`((x: number)=>x+5)(10)`);
+    });
+
+    it('should correcly handle arrow functions without parenthesis around the argument names', () => {
+        const input = `(async x=>x+5)(10)`;
+        expect(typeWiz(input)).toBe(`(async (x: number)=>x+5)(10)`);
+    });
+
     it('should infer `string` type for class method', () => {
         const input = `
             class Greeter {

diff --git a/src/type-collector-snippet.ts b/src/type-collector-snippet.ts
@@ -3,6 +3,7 @@ class NestError extends Error { }
 interface IKey {
     filename: string;
     pos: number;
+    opts: any;
 }
 
 export function getTypeName(value: any, nest = 0): string | null {
@@ -34,8 +35,8 @@ export function getTypeName(value: any, nest = 0): string | null {
 
 const logs: { [key: string]: Set<string> } = {};
 
-export function $_$twiz(name: string, value: any, pos: number, filename: string) {
-    const index = JSON.stringify({ filename, pos } as IKey);
+export function $_$twiz(name: string, value: any, pos: number, filename: string, opts: any) {
+    const index = JSON.stringify({ filename, pos, opts } as IKey);
     try {
         const typeName = getTypeName(value);
         if (!logs[index]) {
@@ -57,8 +58,8 @@ export namespace $_$twiz {
     export const typeName = getTypeName;
     export const get = () => {
         return Object.keys(logs).map((key) => {
-            const { filename, pos } = JSON.parse(key) as IKey;
-            return [filename, pos, Array.from(logs[key])] as [string, number, string[]];
+            const { filename, pos, opts } = JSON.parse(key) as IKey;
+            return [filename, pos, Array.from(logs[key]), opts] as [string, number, string[], any];
         });
     };
 }