Fix: convertFunctionToEs6Class cannot recognize x.prototype = {} pattern

src/services/codefixes/convertFunctionToEs6Class.ts

@@ -57,95 +57,161 @@ namespace ts.codefix {
             const memberElements: ClassElement[] = [];
             // all instance members are stored in the "member" array of symbol
             if (symbol.members) {
-                symbol.members.forEach(member => {
+                symbol.members.forEach((member, key) => {
+                    if (key === "constructor") {
+                        // fn.prototype.constructor = fn
+                        changes.delete(sourceFile, member.valueDeclaration.parent);
+                        return;
+                    }
                     const memberElement = createClassElement(member, /*modifiers*/ undefined);
                     if (memberElement) {
-                        memberElements.push(memberElement);
+                        memberElements.push(...memberElement);
                     }
                 });
             }
 
             // all static members are stored in the "exports" array of symbol
             if (symbol.exports) {
                 symbol.exports.forEach(member => {
-                    const memberElement = createClassElement(member, [createToken(SyntaxKind.StaticKeyword)]);
-                    if (memberElement) {
-                        memberElements.push(memberElement);
+                    if (member.name === "prototype") {
+                        const firstDeclaration = member.declarations[0];
+                        // only one "x.prototype = { ... }" will pass
+                        if (member.declarations.length === 1 &&
+                            isPropertyAccessExpression(firstDeclaration) &&
+                            isBinaryExpression(firstDeclaration.parent) &&
+                            firstDeclaration.parent.operatorToken.kind === SyntaxKind.EqualsToken &&
+                            isObjectLiteralExpression(firstDeclaration.parent.right)
+                        ) {
+                            const prototypes = firstDeclaration.parent.right;
+                            const memberElement = createClassElement(prototypes.symbol, /** modifiers */ undefined);
+                            if (memberElement) {
+                                memberElements.push(...memberElement);
+                            }
+                        }
+                    }
+                    else {
+                        const memberElement = createClassElement(member, [createToken(SyntaxKind.StaticKeyword)]);
+                        if (memberElement) {
+                            memberElements.push(...memberElement);
+                        }
                     }
                 });
             }
 
             return memberElements;
 
-            function shouldConvertDeclaration(_target: PropertyAccessExpression, source: Expression) {
-                // Right now the only thing we can convert are function expressions - other values shouldn't get
-                // transformed. We can update this once ES public class properties are available.
-                return isFunctionLike(source);
+            function shouldConvertDeclaration(_target: PropertyAccessExpression | ObjectLiteralExpression, source: Expression) {
+                // Right now the only thing we can convert are function expressions, get/set accessors and methods
+                // other values like normal value fields ({a: 1}) shouldn't get transformed.
+                // We can update this once ES public class properties are available.
+                if (isPropertyAccessExpression(_target)) {
+                    if (isConstructorAssignment(_target)) return true;
+                    return isFunctionLike(source);
+                }
+                else {
+                    return every(_target.properties, property => {
+                        // a() {}
+                        if (isMethodDeclaration(property) || isGetOrSetAccessorDeclaration(property)) return true;
+                        // a: function() {}
+                        if (isPropertyAssignment(property) && isFunctionExpression(property.initializer) && !!property.name) return true;
+                        // x.prototype.constructor = fn
+                        if (isConstructorAssignment(property)) return true;
+                        return false;
+                    });
+                }
             }
 
-            function createClassElement(symbol: Symbol, modifiers: Modifier[] | undefined): ClassElement | undefined {
+            function createClassElement(symbol: Symbol, modifiers: Modifier[] | undefined): readonly ClassElement[] {
                 // Right now the only thing we can convert are function expressions, which are marked as methods
-                if (!(symbol.flags & SymbolFlags.Method)) {
-                    return;
+                // or { x: y } type prototype assignments, which are marked as ObjectLiteral
+                const members: ClassElement[] = [];
+                if (!(symbol.flags & SymbolFlags.Method) && !(symbol.flags & SymbolFlags.ObjectLiteral)) {
+                    return members;
                 }
 
-                const memberDeclaration = symbol.valueDeclaration as PropertyAccessExpression;
+                const memberDeclaration = symbol.valueDeclaration as PropertyAccessExpression | ObjectLiteralExpression;
                 const assignmentBinaryExpression = memberDeclaration.parent as BinaryExpression;
+                const assignmentExpr = assignmentBinaryExpression.right;
 
-                if (!shouldConvertDeclaration(memberDeclaration, assignmentBinaryExpression.right)) {
-                    return;
+                if (!shouldConvertDeclaration(memberDeclaration, assignmentExpr)) {
+                    return members;
                 }
 
                 // delete the entire statement if this expression is the sole expression to take care of the semicolon at the end
                 const nodeToDelete = assignmentBinaryExpression.parent && assignmentBinaryExpression.parent.kind === SyntaxKind.ExpressionStatement
                     ? assignmentBinaryExpression.parent : assignmentBinaryExpression;
                 changes.delete(sourceFile, nodeToDelete);
 
-                if (!assignmentBinaryExpression.right) {
-                    return createProperty([], modifiers, symbol.name, /*questionToken*/ undefined,
-                        /*type*/ undefined, /*initializer*/ undefined);
+                if (!assignmentExpr) {
+                    members.push(createProperty([], modifiers, symbol.name, /*questionToken*/ undefined,
+                        /*type*/ undefined, /*initializer*/ undefined));
+                    return members;
                 }
 
-                switch (assignmentBinaryExpression.right.kind) {
-                    case SyntaxKind.FunctionExpression: {
-                        const functionExpression = assignmentBinaryExpression.right as FunctionExpression;
-                        const fullModifiers = concatenate(modifiers, getModifierKindFromSource(functionExpression, SyntaxKind.AsyncKeyword));
-                        const method = createMethod(/*decorators*/ undefined, fullModifiers, /*asteriskToken*/ undefined, memberDeclaration.name, /*questionToken*/ undefined,
-                            /*typeParameters*/ undefined, functionExpression.parameters, /*type*/ undefined, functionExpression.body);
-                        copyLeadingComments(assignmentBinaryExpression, method, sourceFile);
-                        return method;
-                    }
+                // f.x = expr
+                if (isPropertyAccessExpression(memberDeclaration) && (isFunctionExpression(assignmentExpr) || isArrowFunction(assignmentExpr))) {
+                    return createFunctionLikeExpressionMember(members, assignmentExpr, memberDeclaration.name);
+                }
+                // f.prototype = { ... }
+                else if (isObjectLiteralExpression(assignmentExpr)) {
+                    return flatMap(
+                        assignmentExpr.properties,
+                        property => {
+                            if (isMethodDeclaration(property) || isGetOrSetAccessorDeclaration(property)) {
+                                // MethodDeclaration and AccessorDeclaration can appear in a class directly
+                                return members.concat(property);
+                            }
+                            if (isPropertyAssignment(property) && isFunctionExpression(property.initializer)) {
+                                return createFunctionLikeExpressionMember(members, property.initializer, property.name);
+                            }
+                            // Drop constructor assignments
+                            if (isConstructorAssignment(property)) return members;
+                            return [];
+                        }
+                    );
+                }
+                else {
+                    // Don't try to declare members in JavaScript files
+                    if (isSourceFileJS(sourceFile)) return members;
+                    if (!isPropertyAccessExpression(memberDeclaration)) return members;
+                    const prop = createProperty(/*decorators*/ undefined, modifiers, memberDeclaration.name, /*questionToken*/ undefined, /*type*/ undefined, assignmentExpr);
+                    copyLeadingComments(assignmentBinaryExpression.parent, prop, sourceFile);
+                    members.push(prop);
+                    return members;
+                }
 
-                    case SyntaxKind.ArrowFunction: {
-                        const arrowFunction = assignmentBinaryExpression.right as ArrowFunction;
-                        const arrowFunctionBody = arrowFunction.body;
-                        let bodyBlock: Block;
+                type MethodName = Parameters<typeof createMethod>[3];
 
-                        // case 1: () => { return [1,2,3] }
-                        if (arrowFunctionBody.kind === SyntaxKind.Block) {
-                            bodyBlock = arrowFunctionBody as Block;
-                        }
-                        // case 2: () => [1,2,3]
-                        else {
-                            bodyBlock = createBlock([createReturn(arrowFunctionBody)]);
-                        }
-                        const fullModifiers = concatenate(modifiers, getModifierKindFromSource(arrowFunction, SyntaxKind.AsyncKeyword));
-                        const method = createMethod(/*decorators*/ undefined, fullModifiers, /*asteriskToken*/ undefined, memberDeclaration.name, /*questionToken*/ undefined,
-                            /*typeParameters*/ undefined, arrowFunction.parameters, /*type*/ undefined, bodyBlock);
-                        copyLeadingComments(assignmentBinaryExpression, method, sourceFile);
-                        return method;
-                    }
+                function createFunctionLikeExpressionMember(members: readonly ClassElement[], expression: FunctionExpression | ArrowFunction, name: MethodName) {
+                    if (isFunctionExpression(expression)) return createFunctionExpressionMember(members, expression, name);
+                    else return createArrowFunctionExpressionMember(members, expression, name);
+                }
 
-                    default: {
-                        // Don't try to declare members in JavaScript files
-                        if (isSourceFileJS(sourceFile)) {
-                            return;
-                        }
-                        const prop = createProperty(/*decorators*/ undefined, modifiers, memberDeclaration.name, /*questionToken*/ undefined,
-                            /*type*/ undefined, assignmentBinaryExpression.right);
-                        copyLeadingComments(assignmentBinaryExpression.parent, prop, sourceFile);
-                        return prop;
+                function createFunctionExpressionMember(members: readonly ClassElement[], functionExpression: FunctionExpression, name: MethodName) {
+                    const fullModifiers = concatenate(modifiers, getModifierKindFromSource(functionExpression, SyntaxKind.AsyncKeyword));
+                    const method = createMethod(/*decorators*/ undefined, fullModifiers, /*asteriskToken*/ undefined, name, /*questionToken*/ undefined,
+    /*typeParameters*/ undefined, functionExpression.parameters, /*type*/ undefined, functionExpression.body);
+                    copyLeadingComments(assignmentBinaryExpression, method, sourceFile);
+                    return members.concat(method);
+                }
+
+                function createArrowFunctionExpressionMember(members: readonly ClassElement[], arrowFunction: ArrowFunction, name: MethodName) {
+                    const arrowFunctionBody = arrowFunction.body;
+                    let bodyBlock: Block;
+
+                    // case 1: () => { return [1,2,3] }
+                    if (arrowFunctionBody.kind === SyntaxKind.Block) {
+                        bodyBlock = arrowFunctionBody as Block;
                     }
+                    // case 2: () => [1,2,3]
+                    else {
+                        bodyBlock = createBlock([createReturn(arrowFunctionBody)]);
+                    }
+                    const fullModifiers = concatenate(modifiers, getModifierKindFromSource(arrowFunction, SyntaxKind.AsyncKeyword));
+                    const method = createMethod(/*decorators*/ undefined, fullModifiers, /*asteriskToken*/ undefined, name, /*questionToken*/ undefined,
+                        /*typeParameters*/ undefined, arrowFunction.parameters, /*type*/ undefined, bodyBlock);
+                    copyLeadingComments(assignmentBinaryExpression, method, sourceFile);
+                    return members.concat(method);
                 }
             }
         }
@@ -189,4 +255,10 @@ namespace ts.codefix {
     function getModifierKindFromSource(source: Node, kind: SyntaxKind): readonly Modifier[] | undefined {
         return filter(source.modifiers, modifier => modifier.kind === kind);
     }
+
+    function isConstructorAssignment(x: ObjectLiteralElementLike | PropertyAccessExpression) {
+        if (!x.name) return false;
+        if (isIdentifier(x.name) && x.name.text === "constructor") return true;
+        return false;
+    }
 }

tests/cases/fourslash/convertFunctionToEs6Class-propertyMember.ts

@@ -0,0 +1,31 @@
+// @allowNonTsExtensions: true
+// @Filename: test123.js
+
+/// <reference path="./fourslash.ts" />
+
+//// // Comment
+//// function /*1*/fn() {
+////     this.baz = 10;
+//// }
+//// fn.prototype = {
+////     /** JSDoc */
+////     bar: function () {
+////         console.log('hello world');
+////     }
+//// }
+
+verify.codeFix({
+    description: "Convert function to an ES2015 class",
+    newFileContent:
+        `// Comment
+class fn {
+    constructor() {
+        this.baz = 10;
+    }
+    /** JSDoc */
+    bar() {
+        console.log('hello world');
+    }
+}
+`,
+});

tests/cases/fourslash/convertFunctionToEs6Class-prototypeAccessor.ts

@@ -0,0 +1,31 @@
+// @allowNonTsExtensions: true
+// @Filename: test123.js
+
+/// <reference path="./fourslash.ts" />
+
+//// // Comment
+//// function /*1*/fn() {
+////     this.baz = 10;
+//// }
+//// fn.prototype = {
+////     /** JSDoc */
+////     get bar() {
+////         return this.baz;
+////     }
+//// }
+
+verify.codeFix({
+    description: "Convert function to an ES2015 class",
+    newFileContent:
+        `// Comment
+class fn {
+    constructor() {
+        this.baz = 10;
+    }
+    /** JSDoc */
+    get bar() {
+        return this.baz;
+    }
+}
+`,
+});

tests/cases/fourslash/convertFunctionToEs6Class-prototypeMethod.ts

@@ -0,0 +1,29 @@
+// @allowNonTsExtensions: true
+// @Filename: test123.js
+
+/// <reference path="./fourslash.ts" />
+
+//// // Comment
+//// function /*1*/fn() {
+////     this.baz = 10;
+//// }
+//// fn.prototype = {
+////     bar() {
+////         console.log('hello world');
+////     }
+//// }
+
+verify.codeFix({
+    description: "Convert function to an ES2015 class",
+    newFileContent:
+        `// Comment
+class fn {
+    constructor() {
+        this.baz = 10;
+    }
+    bar() {
+        console.log('hello world');
+    }
+}
+`,
+});

tests/cases/fourslash/convertFunctionToEs6Class-removeConstructor.ts

@@ -0,0 +1,24 @@
+// @allowNonTsExtensions: true
+// @Filename: test123.js
+
+/// <reference path="./fourslash.ts" />
+
+//// // Comment
+//// function /*1*/fn() {
+////     this.baz = 10;
+//// }
+//// fn.prototype = {
+////     constructor: fn
+//// }
+
+verify.codeFix({
+    description: "Convert function to an ES2015 class",
+    newFileContent:
+        `// Comment
+class fn {
+    constructor() {
+        this.baz = 10;
+    }
+}
+`,
+});

tests/cases/fourslash/convertFunctionToEs6Class-removeConstructor2.ts

@@ -0,0 +1,22 @@
+// @allowNonTsExtensions: true
+// @Filename: test123.js
+
+/// <reference path="./fourslash.ts" />
+
+//// // Comment
+//// function /*1*/fn() {
+////     this.baz = 10;
+//// }
+//// fn.prototype.constructor = fn
+
+verify.codeFix({
+    description: "Convert function to an ES2015 class",
+    newFileContent:
+        `// Comment
+class fn {
+    constructor() {
+        this.baz = 10;
+    }
+}
+`,
+});